mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 14:21:06 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
fc4edd2ca5 | ||
|
|
de896a8e96 | ||
|
|
b038240369 | ||
|
|
9abc623f2c | ||
|
|
3c38c22649 | ||
|
|
18b9c18454 | ||
|
|
b317c34698 | ||
|
|
455fc8a794 | ||
|
|
3af82092d6 | ||
|
|
a9e9eb2ae1 | ||
|
|
5a7f0860e7 | ||
|
|
e52051dc71 | ||
|
|
82cf9b83c8 | ||
|
|
2aa6ace036 | ||
|
|
fc21e54214 | ||
|
|
9d37a3ad5f | ||
|
|
60445b58c0 | ||
|
|
d05d26e9c3 | ||
|
|
34d4178952 | ||
|
|
30e5ba4a31 | ||
|
|
6f66917530 | ||
|
|
47de4a4d83 | ||
|
|
30875ed297 | ||
|
|
31da8f5dca | ||
|
|
4ba7c94d4d | ||
|
|
c472ff04c8 | ||
|
|
512ae42f7c | ||
|
|
5b734e427c | ||
|
|
82d7e82926 | ||
|
|
772c9330e9 | ||
|
|
6f8536aed1 | ||
|
|
6cc16bee09 | ||
|
|
95f2860231 | ||
|
|
22ee8d0e9a | ||
|
|
0e4c3e8727 | ||
|
|
cd49a498c8 | ||
|
|
ce2d83b62d | ||
|
|
42bc3dd9ad | ||
|
|
a2edcb12f9 | ||
|
|
8896607268 |
No files matched your search
+34
-3
@@ -8,7 +8,7 @@
|
||||
# below refuses anything older: silently building for a pre-Ampere 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.34 LANGUAGES CXX)
|
||||
project(strata VERSION 0.1.37 LANGUAGES CXX)
|
||||
add_compile_definitions(STRATA_VERSION="${PROJECT_VERSION}")
|
||||
# issue #31: release builds keep a PDB beside the exe (not shipped in it), so a minidump from a stalled engine can
|
||||
# be read back to the exact source line
|
||||
@@ -350,6 +350,14 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
|
||||
add_test(NAME file_expert_source_test COMMAND file_expert_source_test)
|
||||
set_tests_properties(file_expert_source_test PROPERTIES TIMEOUT 60)
|
||||
endif()
|
||||
# #477: the learned expert profile (--expert-profile-save): its order and make_profile.py's file format (CPU only)
|
||||
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/core/expert_profile_save_test.cpp")
|
||||
enable_testing()
|
||||
add_executable(expert_profile_save_test tests/core/expert_profile_save_test.cpp)
|
||||
target_link_libraries(expert_profile_save_test PRIVATE strata_engine)
|
||||
add_test(NAME expert_profile_save_test COMMAND expert_profile_save_test)
|
||||
set_tests_properties(expert_profile_save_test PROPERTIES TIMEOUT 60)
|
||||
endif()
|
||||
if(STRATA_BUILD_TESTS AND STRATA_NATIVE_EXPERTS)
|
||||
# native_experts.txt v2-v4 (the real packs' copies and the per-role shard column) and the arena filled from
|
||||
# synthetic GGUF shards at the real expert dimensions
|
||||
@@ -580,6 +588,14 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
|
||||
set_tests_properties(sampler_parity_one_block PROPERTIES ENVIRONMENT "STRATA_OLD_SAMPLER=0;STRATA_SAMPLER_ONE_BLOCK=1")
|
||||
add_test(NAME sampler_parity_old COMMAND sampler_parity --selftest)
|
||||
set_tests_properties(sampler_parity_old PROPERTIES ENVIRONMENT "STRATA_OLD_SAMPLER=1;STRATA_SAMPLER_ONE_BLOCK=0")
|
||||
# S19: the decode QSA top-k and the greedy argmax on thread-block clusters (sm_90+) vs the one-CTA kernels, bitwise,
|
||||
# captured in a graph too (GPU, synthetic, no model; skips below sm_90); --bench: per-call times
|
||||
if(STRATA_ENABLE_CUDA)
|
||||
add_executable(decode_cluster_parity src/kernels/decode_cluster_parity.cpp)
|
||||
target_include_directories(decode_cluster_parity PRIVATE ${_strata_gpu_include_directories})
|
||||
target_link_libraries(decode_cluster_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target})
|
||||
add_test(NAME decode_cluster_parity COMMAND decode_cluster_parity --selftest)
|
||||
endif()
|
||||
|
||||
add_executable(rope_parity src/kernels/rope_parity.cpp)
|
||||
target_include_directories(rope_parity PRIVATE ${_strata_gpu_include_directories})
|
||||
@@ -868,7 +884,10 @@ if(STRATA_NATIVE_EXPERTS)
|
||||
set(_strata_gg ${strata_llamacpp_SOURCE_DIR}/ggml)
|
||||
endif()
|
||||
set(_strata_mmq ${_strata_gg}/src/ggml-cuda)
|
||||
add_library(strata_mmq STATIC src/prefill/moe_mmq.cu src/prefill/ggml_cuda_host.cu ${_strata_mmq}/quantize.cu)
|
||||
# #136: the Q2_0 pack's prompt experts on Strata's own int8 kernels (moe_fused.cu, opt-in STRATA_PF_FUSED=1; CUDA
|
||||
# only: the HIP branch below keeps MMQ, and an sm_75 device or an sm_75-only build reports them unavailable)
|
||||
add_library(strata_mmq STATIC src/prefill/moe_mmq.cu src/prefill/moe_fused.cu src/prefill/moe_fused_iq.cu
|
||||
src/prefill/ggml_cuda_host.cu ${_strata_mmq}/quantize.cu)
|
||||
# q8_0: the drafter's matrices (E-9); STRATA_ORCA_Q4KS_MMQ adds q5_0 (#296, in _strata_mmq_types), and
|
||||
# STRATA_MMQ_KQUANTS q4_k q5_k q5_1, the experts of Unsloth's UD-Q4_K_XL (docs/UNSLOTH_Q4.md; the dispatch
|
||||
# follows eddoursul/Strata 7706c1b)
|
||||
@@ -893,7 +912,19 @@ if(STRATA_NATIVE_EXPERTS)
|
||||
target_compile_definitions(strata_mmq PRIVATE STRATA_ORCA_Q4KS_MMQ=1)
|
||||
endif()
|
||||
target_link_libraries(strata_prefill PUBLIC strata_mmq)
|
||||
target_compile_definitions(strata_prefill PUBLIC STRATA_PREFILL_MMQ=1)
|
||||
target_compile_definitions(strata_prefill PUBLIC STRATA_PREFILL_MMQ=1 STRATA_PREFILL_FUSED=1)
|
||||
if(STRATA_BUILD_TESTS)
|
||||
# #136: the fused int8 experts against a double-precision reference and the MMQ path (exit 77 below sm_80)
|
||||
add_executable(prefill_fused_moe_test tests/cuda/prefill_fused_moe_test.cpp)
|
||||
target_link_libraries(prefill_fused_moe_test PRIVATE strata_mmq ggml-base ${_strata_gpu_runtime_target})
|
||||
add_test(NAME prefill_fused_moe_test COMMAND prefill_fused_moe_test)
|
||||
set_tests_properties(prefill_fused_moe_test PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 300)
|
||||
# #136: the same for the native packs' formats (IQ2_XXS IQ2_XS IQ2_S IQ3_XXS IQ3_S IQ4_XS, down Q2_0 IQ4_NL)
|
||||
add_executable(prefill_fused_iq_test tests/cuda/prefill_fused_iq_test.cpp)
|
||||
target_link_libraries(prefill_fused_iq_test PRIVATE strata_mmq ggml-base ${_strata_gpu_runtime_target})
|
||||
add_test(NAME prefill_fused_iq_test COMMAND prefill_fused_iq_test)
|
||||
set_tests_properties(prefill_fused_iq_test PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 600)
|
||||
endif()
|
||||
if(STRATA_BUILD_TESTS AND STRATA_MMQ_KQUANTS)
|
||||
# the UD-Q4_K_XL expert formats through MMQ against ggml's dequantized weights (synthetic, a few seconds)
|
||||
add_executable(prefill_mmq_kquant_test tests/cuda/prefill_mmq_kquant_test.cpp)
|
||||
|
||||
@@ -23,7 +23,7 @@ word, so 60 tokens per second is faster than you can read.
|
||||
|
||||
| Size | Writes answers | Reads your prompt |
|
||||
| --- | ---: | ---: |
|
||||
| **Q2_0** | 93 tokens/s | 2,170 tokens/s |
|
||||
| **Q2_0** | 94 tokens/s | 2,650 tokens/s |
|
||||
| **IQ2_XS** | 79 tokens/s | 2,090 tokens/s |
|
||||
| **IQ3_XXS** | 62 tokens/s | 1,750 tokens/s |
|
||||
| **IQ3_S** | 53 tokens/s | 1,620 tokens/s |
|
||||
@@ -87,7 +87,8 @@ it left off) and starts it. Your browser opens the Strata app at `http://127.0.0
|
||||
> window. The window tells you what it is doing.
|
||||
|
||||
**Next time**, run `START-HERE.bat` (or `./setup.sh`) again: it starts right away, nothing is downloaded twice. Close
|
||||
its window to stop the model. Updating, Docker, several cards, where the files go and every option:
|
||||
its window to stop the model. `UPDATE.bat` (`./update.sh`) updates Strata without starting it. Updating, Docker,
|
||||
several cards, where the files go and every option:
|
||||
[docs/INSTALL.md](docs/INSTALL.md).
|
||||
|
||||
## Which model should I pick?
|
||||
@@ -103,7 +104,8 @@ is faster, larger is a bit smarter.
|
||||
| **96 GB or more** | **IQ3_S**, or Unsloth's 4-bit (experimental) | room for the largest sizes with everything else open |
|
||||
|
||||
- **[Coder](docs/MODELS.md#coder)** - a coding version with half of the experts removed: 91% of the full model's
|
||||
SWE-bench Verified score (by its authors), fits 32 GB of RAM. Weaker outside coding.
|
||||
SWE-bench Verified score (by its authors), fits 32 GB of RAM. Weaker outside code, including Chinese and other
|
||||
CJK text (#438): for those, take Q2_0, IQ2_XS or IQ3_S, which keep every expert.
|
||||
- **[Swift 1.5](docs/MODELS.md#swift-15)** - a fine-tune that thinks much shorter before it answers, so you get the
|
||||
answer sooner, at about the same quality.
|
||||
- **[Unsloth UD-Q4_K_XL](docs/MODELS.md#unsloth-ud-q4_k_xl-experimental)** (experimental) - the closest to the full
|
||||
|
||||
+36
@@ -0,0 +1,36 @@
|
||||
@echo off
|
||||
rem Update Strata without starting the model (#475): the newest code (git pull, when this folder is a git clone), then
|
||||
rem what START-HERE.bat does before a start - the engine (a new one when this version needs it), the Python packages,
|
||||
rem each installed model's settings and draft subset. The model files are not touched. Start the model later with
|
||||
rem START-HERE.bat. Options are passed on to setup.py.
|
||||
setlocal
|
||||
title Strata - update
|
||||
cd /d "%~dp0"
|
||||
rem All of it in one block: cmd reads a .bat file while it runs it, and the git pull can change this very file.
|
||||
(
|
||||
if exist ".git" (
|
||||
where git >nul 2>nul || (
|
||||
echo This folder is a git clone, but git is not on PATH: install Git for Windows ^(winget install Git.Git^)
|
||||
echo or run "git pull" here yourself, then run UPDATE.bat again.
|
||||
pause
|
||||
exit /b 1
|
||||
)
|
||||
echo Getting the newest Strata ^(git pull^) ...
|
||||
git pull --ff-only || (
|
||||
echo.
|
||||
echo git pull did not succeed ^(the reason is above^): nothing was updated. Files you changed here can stop it:
|
||||
echo "git status" lists them.
|
||||
pause
|
||||
exit /b 1
|
||||
)
|
||||
) else (
|
||||
echo This copy of Strata was not made with git, so it cannot fetch new files itself. Download the newest one:
|
||||
echo https://github.com/Niko1221/Strata/archive/refs/heads/main.zip
|
||||
echo unzip it anywhere and run START-HERE.bat ^(or UPDATE.bat^) in it: it finds the model files in Strata-data
|
||||
echo and sets itself up the same way - nothing big is downloaded again.
|
||||
echo Checking this copy's engine and settings meanwhile ...
|
||||
)
|
||||
call "%~dp0START-HERE.bat" --update %*
|
||||
if not errorlevel 1 pause
|
||||
exit /b
|
||||
)
|
||||
@@ -68,6 +68,11 @@ On a PC with no NVIDIA card Strata can use, the AMD card is chosen by itself; wi
|
||||
amd_comgr and the Microsoft C++ runtime; ROCm 10.2.0a20260930 from AMD's TheRock builds, licenses in
|
||||
`engine\rocm\licenses`). The HIP runtime works through the AMD driver's own components, so the driver is the one
|
||||
thing it needs from the PC.
|
||||
- **The HIP runtime next to `strata.exe` (0.1.35, #468 #461):** `amdhip64_7.dll` and `amd_comgr.dll` are also put in
|
||||
`engine\` (setup copies them there on every start). Windows looks in the program's folder before System32, where
|
||||
some AMD drivers install their own `amdhip64_7.dll`; with that one, the bundled libraries crashed on the first
|
||||
prompt (an access violation, or `hipErrorInvalidDeviceFunction`). The engine's log names the runtime it loaded
|
||||
(`strata generate: HIP runtime ...`).
|
||||
- **Before the ~60 GB model download** setup runs `engine\strata-device.exe --list-devices` (with `engine\rocm\bin` on
|
||||
the PATH): if the HIP runtime does not see the card, setup stops there and points to the driver. It also gives the
|
||||
card's HIP number: with an integrated Radeon that is device 1, not 0 (#325). From then on setup lists the AMD cards
|
||||
@@ -79,6 +84,13 @@ On a PC with no NVIDIA card Strata can use, the AMD card is chosen by itself; wi
|
||||
succeeded, on Windows only. `hipHostGetDevicePointer` returns the host pointer itself on Windows: kernels read
|
||||
mapped memory through it correctly, but a device-to-device copy into it does not land, which is why
|
||||
`tests/hip/handoff` times out there (the engine does not use that copy; `tests/hip/mapped_alias` reports it).
|
||||
- Two more (#380, #377, by BlueKingMuch, measured on an RX 6800 that drives the desktop): `hipMemGetInfo` on Windows
|
||||
does not subtract what the desktop and other programs hold on the card, so `--expert-cache auto` filled the card
|
||||
past what Windows keeps in VRAM and decode fell from 41 to 30 tok/s. The engine now lowers that free figure by
|
||||
what Windows' video memory budget for the process withholds (logged once: `strata: Windows budgets N of this
|
||||
card's M MiB ...`; `STRATA_WDDM_BUDGET=0` turns it off). And the PCIe probe times its copies on the host clock
|
||||
there, since HIP's events read impossible speeds (3,300-26,000 GB/s), so a slow link now gets a smaller
|
||||
`pcie_frac` as on NVIDIA.
|
||||
|
||||
**What is validated (0.1.34):** #325's author ran the engine of this port on an RX 9070 XT (Windows 11, ROCm
|
||||
10.2.0a20260930 in `.venv`, compiled on the PC): Coder IQ1_M at 32K, 29.2 tok/s decode, ~181 tok/s prefill,
|
||||
|
||||
+27
-2
@@ -19,6 +19,15 @@ RTX 5070 **12 GB**, Ryzen 5 7600 (6 cores), 64 GB DDR5-5200, Windows, engine 0.1
|
||||
MTP speculative decoding on. "262K" is the model's full context window (a 259,943-token prompt). The IQ2_XS row was
|
||||
measured with Swift 1.5's IQ2_XS, which runs at the original's speed.
|
||||
|
||||
**Engine 0.1.36 (#136), the same PC:** Q2_0's prompt experts run on fused int8 tensor-core kernels (RTX 30 and newer):
|
||||
4K 1,294 -> 1,570, 32K 2,170 -> 2,653, 128K 2,123 -> 2,468 tokens/s (+16-22%), as close to an FP16 reference as the
|
||||
previous kernels (closer at 32K: teacher-forced KL 0.009 vs 0.012). The decode path's block selection and greedy
|
||||
argmax run on thread-block clusters (RTX 50, sm_90+; other cards keep the previous kernels; the same tokens): Q2_0 output at 4K 89 -> 93.5, at 128K
|
||||
64.5 -> 76.4 tokens/s. `STRATA_PF_FUSED=0` keeps the previous prompt kernels (byte-identical answers to 0.1.35);
|
||||
`STRATA_PF_FUSED=1` also runs the native IQ packs' fused kernels (opt-in: IQ2_XS prompts +12% at 4K, +3% at 32K, the
|
||||
IQ3 packs about even); `STRATA_QSA_CLUSTER=0` / `STRATA_ARGMAX_MULTI=0` turn the decode kernels off. The tables
|
||||
below are 0.1.26's.
|
||||
|
||||
### Prompt processing (tokens/s)
|
||||
|
||||
| Model | 1K | 4K | 32K | 64K | 128K | 262K |
|
||||
@@ -98,7 +107,10 @@ English/code subset from before (40,525 ids, ~110 MiB less VRAM, English answers
|
||||
almost no drafts). `--draft-vocab cyrillic` takes the English/code subset plus the whole Cyrillic script (58,963
|
||||
ids): the shipped subsets hold 142 of the vocabulary's 18,580 Cyrillic tokens, so Ukrainian or Russian answers got
|
||||
1.4 tokens a round; with it 2.1, and 83 -> 109 tokens/s (RTX 5090, the NVFP4 fork), English unchanged.
|
||||
`tools/draft_vocab.py` builds and inspects subsets.
|
||||
`tools/draft_vocab.py` builds and inspects subsets. When the start stops with "the draft head does not fit" (a
|
||||
12 GB card with a long context, #474), the engine says how much the head needs, how much VRAM is free and which
|
||||
smaller subset fits, and the server's start error repeats it; setup suggests `--draft-vocab en` on cards under
|
||||
14 GB (only a suggestion: nothing changes unless you pass it).
|
||||
|
||||
**Low-RAM mode (engine 0.1.26, chosen by setup):** normally all of a model's experts are copied into RAM (23-50 GB,
|
||||
pinned) and the GPU holds a copy of the most-used ones. On a PC whose RAM cannot hold them beside the system (the
|
||||
@@ -159,7 +171,8 @@ RTX 5070, against ~3 tokens/s before these changes.
|
||||
RAM copy, blobs and MB from the files, the time spent reading them; `routing prefetch`: how many of the file reads had
|
||||
been warmed). The server log has the same per request (`expert tiers: GPU ... hits ...; RAM ... blobs, files ...
|
||||
blobs ... MB read`), and `GET /metrics` lists `ram_blobs`, `file_blobs` and `file_mb` for each recent request (with
|
||||
engine 0.1.31 or newer).
|
||||
engine 0.1.31 or newer). It also lists each request's speculative drafts, `drafts_offered` and `drafts_accepted`
|
||||
(`null` when the engine did not report them), and their sums since the server started in `totals` (#457).
|
||||
|
||||
Time to first token is prompt length / prompt speed: with Q2_0 about 4 s at 4K, 25 s at 32K, under 2 minutes at 128K
|
||||
and 4.5 minutes at 262K (engine 0.1.13 made long prompts about twice as fast, below).
|
||||
@@ -400,6 +413,17 @@ the OS file cache, so loading again takes seconds while that RAM is not needed e
|
||||
16 GB with Q2_0 in the low-RAM mode: unloading takes ~0.3 s, and a request to an unloaded model answered after
|
||||
4.6 s (text) or 14.7 s (a picture, image encoder on the CPU).
|
||||
|
||||
**Keep what the expert cache learned across restarts (opt-in, engine 0.1.36, #477):** a start fills the GPU's expert
|
||||
cache from the shipped profile, and the adaptive tier (`--adapt-every`) then moves in the experts your requests use.
|
||||
With `"expert_profile_save": "expert-profile-learned.bin"` in `strata-<model>.json` the engine saves that as a
|
||||
profile - the experts in VRAM first, then the routing it counted since the start, then the shipped order - on a
|
||||
clean exit and every 10 minutes between requests (`"expert_profile_save_every": 5` for another interval, `0` for
|
||||
exit only), written to a temporary file and renamed, so a crash never leaves half a file. The next start begins from
|
||||
it instead of the config's `--expert-profile` when it is a profile of the same model (else from the config's, as
|
||||
before). A relative path is in the Strata folder; one file per model, and a profile per project works the same way
|
||||
(point the key at another file). The file is a fingerprint of what you used the model for: it stays on your PC.
|
||||
Without the key nothing is counted or written. Setup rewrites the config when run again: add the key again then.
|
||||
|
||||
---
|
||||
|
||||
## Using it
|
||||
@@ -792,6 +816,7 @@ the document, +0.4% on the chat. Details: `bench/results/2026-09-27-esp/`.
|
||||
| Pictures are slow (10-30 s) | The encoder runs on the CPU: run setup again with `--vision gpu` (needs ~1.4 GB of VRAM). |
|
||||
| A request never finishes: "reading the prompt", GPU "100%" at low power | The GPU ran out of VRAM (engines before 0.1.9 could end with ~30 MiB free at large contexts). Run `START-HERE.bat` once to get engine 0.1.9 or newer; the log then says `... MiB of VRAM free with everything loaded` (a few hundred) and names the `--vram-reserve-mib` to add if it is low. |
|
||||
| Generation stops mid-answer, GPU "100%", one CPU core busy | Fixed in engine 0.1.12 (issue #29, a race in the CPU expert pool on big-VRAM cards). Since then a request that stops moving ends with an error instead of hanging (after 2 minutes; 1 minute from 0.1.13): the log says `no progress for ... s ... (issue #29)` with where it stopped, and the next request starts the engine again. If you see that line, please open an issue with it. Engine 0.1.13 adds a stall report under it (what every expert-pool thread and the GPU handshake were doing, memory and page faults) and, on Windows, a `strata-stall-<pid>.dmp` file with every thread's stack: attach both. (`STRATA_WATCHDOG_S` sets the time in seconds; 0 turns it off.) Engine 0.1.14 fixes the stall those reports found (issue #31: with the IQ packs the host could wait forever inside the NVIDIA driver while copying experts in a verify window; the experts are now copied by a GPU kernel, `--pcie-mode dma` restores the old way). |
|
||||
| `the engine said nothing for ... s during the request` or `... did not finish the request after it was stopped (STOP)` | Issue #481: the engine and the server lost step (the engine waits for its next command, the server for the request's end; GPU at 0 %, nothing in the log). The server ends the engine after 300 s without a line from it during a request (while a prompt is read: each chunk may take three times the previous one's time, the first one up to its tokens at 50 tok/s more), the request ends with an error and the next request starts the engine again. `"engine_silence_s": 600` in `strata-<model>.json` sets the time (0 = wait forever, as before). If you see it, please add the end of the engine log to #481. |
|
||||
| `out of memory: cudaFuncSetAttribute` in the log (IQ3_XXS, long prompt) | Fixed in engine 0.1.15: CUDA loaded a kernel's code when it was first needed, and mid-prompt there was no VRAM left for it. Run `START-HERE.bat` (Windows) or `./setup.sh` (Linux) once to update. |
|
||||
| Anything else | The engine log is `strata-<model>.log` in this folder. |
|
||||
|
||||
|
||||
+29
-5
@@ -123,11 +123,20 @@ The same idea, in a container (NVIDIA cards).
|
||||
|
||||
## Updating
|
||||
|
||||
Download the new version and unzip it anywhere (or `git pull`), then run `START-HERE.bat` (Linux: `./setup.sh`) in
|
||||
it. The model files are kept in a `Strata-data` folder next to your Strata folder, so a new copy finds them and sets
|
||||
itself up the same way - nothing big is downloaded again. On Linux after a `git pull`, setup compiles the engine
|
||||
again when its source changed (a minute or two for the changed files); if that compile fails, it says so and starts
|
||||
the engine you had.
|
||||
**`UPDATE.bat`** (Linux: `./update.sh`) updates Strata without starting the model - for when the GPU is busy with
|
||||
something else, or you just want the new version ready. In a `git clone` it runs `git pull`, then does what
|
||||
`START-HERE.bat` does before a start: the engine (a new ready-made one when the new version needs it; on Linux a
|
||||
compiled engine is compiled again when its source changed), the Python packages, and each installed model's settings
|
||||
and draft subset. The model files are not touched (at most a new engine is downloaded) and no question is asked.
|
||||
Close the model's window first (a running engine cannot be replaced); start the model later with `START-HERE.bat` as
|
||||
usual. In a copy that was downloaded as a zip it says to download the new zip (below): it cannot fetch new files
|
||||
itself.
|
||||
|
||||
Or by hand: download the new version and unzip it anywhere (or `git pull`), then run `START-HERE.bat` (Linux:
|
||||
`./setup.sh`) in it. The model files are kept in a `Strata-data` folder next to your Strata folder, so a new copy
|
||||
finds them and sets itself up the same way - nothing big is downloaded again. On Linux after a `git pull`, setup
|
||||
compiles the engine again when its source changed (a minute or two for the changed files); if that compile fails, it
|
||||
says so and starts the engine you had.
|
||||
|
||||
## Where things are stored
|
||||
|
||||
@@ -177,11 +186,26 @@ START-HERE.bat --data-dir E:\Strata-data keep the model files somewhere
|
||||
START-HERE.bat --port 8081 another port
|
||||
START-HERE.bat --gpu 1 another GPU (setup picks the one with the most VRAM)
|
||||
START-HERE.bat --gpus 0,2 several GPUs sharing the model
|
||||
START-HERE.bat --vram-reserve-mib 2048 leave 2 GB of VRAM free for other programs (remembered)
|
||||
START-HERE.bat --setup --backend hip the AMD engine on a PC that also has an NVIDIA card
|
||||
START-HERE.bat --setup --host 0.0.0.0 --api-key <secret> reachable from other devices, with a key
|
||||
START-HERE.bat --calibrate tune the engine for this PC (about 5-10 minutes), then start
|
||||
START-HERE.bat --check only check this PC
|
||||
```
|
||||
|
||||
**Leaving VRAM for other programs (#493):** Strata fills the graphics card's free VRAM with experts (the expert cache)
|
||||
and leaves `--vram-reserve-mib` MiB free: 700 by default. For a game, a 3D program or another model beside it, leave
|
||||
more: `START-HERE.bat --vram-reserve-mib 2048` (Linux: `./setup.sh --vram-reserve-mib 2048`) writes it into the
|
||||
model's `strata-<model>.json` and starts it; at setup (`--setup --vram-reserve-mib 2048`) it goes into the new config.
|
||||
By hand: add `"--vram-reserve-mib", "2048"` to the config's `"args"` list and restart. The expert cache is then that
|
||||
much smaller, so answers can be a little slower. The engine sizes its cache from the VRAM free when it starts: what
|
||||
another program already holds then is left alone anyway; the reserve is room for what it takes later.
|
||||
|
||||
**Model files downloaded by hand, or from a mirror (#495):** setup's step 5 prints the folder it expects them in
|
||||
(`Strata-data\models\<SIZE>\`, e.g. `Strata-data\models\IQ3_XXS\`): put them there with their original names, or
|
||||
point setup at them with `--gguf-dir`. To let setup download from a Hugging Face mirror itself, set `HF_ENDPOINT`
|
||||
first (Windows: `set HF_ENDPOINT=https://hf-mirror.com`, Linux: `export HF_ENDPOINT=https://hf-mirror.com`): the same
|
||||
pinned revisions and checks apply, and the MTP draft layer comes from there too.
|
||||
|
||||
On Linux the same options go to `./setup.sh`. `START-HERE.bat --help` lists them all. The server's own settings
|
||||
(sharing the GPU with games, MCP tools, CORS, API keys, the API itself) are in the [details](DETAILS.md#using-it).
|
||||
+5
-2
@@ -24,7 +24,7 @@ Measured on an RTX 5070 (12 GB), a Ryzen 5 7600 and 64 GB of RAM:
|
||||
|
||||
| Size | Writes answers (short chat) | Writes answers (128K context) | Reads your prompt |
|
||||
| --- | ---: | ---: | ---: |
|
||||
| **Q2_0** | 93 tokens/s | 74 tokens/s | 2,170 tokens/s |
|
||||
| **Q2_0** | 94 tokens/s | 76 tokens/s | 2,650 tokens/s |
|
||||
| **IQ2_XS** | 79 tokens/s | 63 tokens/s | 2,090 tokens/s |
|
||||
| **IQ3_XXS** | 62 tokens/s | 49 tokens/s | 1,750 tokens/s |
|
||||
| **IQ3_S** | 53 tokens/s | 46 tokens/s | 1,620 tokens/s |
|
||||
@@ -94,7 +94,10 @@ The original, in all four sizes. With images, and with the
|
||||
half of the experts removed, keeping the ones that code, tool use and images need (91% of the full model's SWE-bench
|
||||
Verified score, 99% of LiveCodeBench, by its authors). One size (IQ1_M: its experts stored like IQ3_S): shard 1 is
|
||||
**29.6 GB**, so it fits a PC with **32 GB of RAM**, runs 262K context on 64 GB, and reads long prompts the fastest of
|
||||
all. Weaker outside coding. More: [details](DETAILS.md#or-the-coder-half-the-experts-for-code).
|
||||
all. It is weaker outside code, and that includes Chinese and other CJK text (#438: Chinese answers came out wrong
|
||||
or looping where English was fine). For general chat or CJK text, take a size that keeps every expert: Q2_0,
|
||||
IQ2_XS or IQ3_S.
|
||||
More: [details](DETAILS.md#or-the-coder-half-the-experts-for-code).
|
||||
|
||||
```
|
||||
START-HERE.bat --setup --family coder
|
||||
|
||||
+2
-1
@@ -8,7 +8,8 @@ speed comes from (decode then barely touches the CPU pool).
|
||||
|
||||
This is pipeline (layer) parallelism, not tensor parallelism: a token crosses from one card to the next once per
|
||||
verify window (a few hundred KB through pinned RAM), not twice per layer. No NVLink or peer-to-peer access is
|
||||
needed; cards on x4 or x1 slots work, and the PCIe share of each card is probed on its own link.
|
||||
needed; cards on x4 or x1 slots work, and the PCIe share of each card is probed on its own link. A `--pcie-frac` you give
|
||||
is every card's share and skips those probes; there is no per-card setting yet.
|
||||
|
||||
## Using it
|
||||
|
||||
|
||||
+10
-2
@@ -25,8 +25,16 @@ What setup does differently for this model:
|
||||
- It needs 48 GB of RAM or more (with less it asks, default no; `--model UD-Q4_K_XL --yes` installs it anyway) and
|
||||
engine 0.1.32 or newer (checked before anything is downloaded), and an NVIDIA
|
||||
GPU: it has not been run on AMD cards (its prompt kernels for the Q4_K / Q5_K experts are NVIDIA-only), so with
|
||||
`--backend hip` setup says so and asks before the download (#429; `--model UD-Q4_K_XL --yes` tries it). One GPU only: with `--gpus` it uses the first one and says so. No images (the vision encoder is not wired to
|
||||
this file yet) and no experimental speed projection (not tested with it).
|
||||
`--backend hip` setup says so and asks before the download (#429; `--model UD-Q4_K_XL --yes` tries it). One GPU
|
||||
by default: the RAM budget below has no layer split (the engine refuses `--resident-budget-gib` with one). No
|
||||
images (the vision encoder is not wired to this file yet) and no experimental speed projection (not tested with it).
|
||||
- Several GPUs (#498): when the RAM holds the GGUF files and 24 GB more (~135 GB of RAM) and two or more cards can
|
||||
share it, setup asks (one GPU stays the default; `--gpus 0,1` takes the split). The split runs **without** the RAM budget: all 77 GB of experts are loaded into RAM from the GGUFs at
|
||||
start, the files pass through the OS file cache while they load, and the config gets `"gpu": [0, 1]` and
|
||||
`"layer_split": "auto"`. Measured on 2x RTX 3090 with 165 GiB (#498): decode 31 tok/s on one card with the budget,
|
||||
64-78 tok/s split (55 tok/s at a 128K prompt), with `MemAvailable` never under 68 GiB. With less RAM, `--gpus`
|
||||
keeps one GPU and says so, and `START-HERE.bat --gpus 0,1` on an installed UD-Q4_K_XL stops with the reason
|
||||
(it used to keep the budget, and the engine exited with code 2).
|
||||
- It downloads the four shards below from the pinned revision `38bb39e` (resumable, like the other models), then
|
||||
checks each one's size and SHA-256 against the table below; the check takes a few minutes once and is remembered
|
||||
in the file's finish mark. A file with the wrong hash is deleted, so the next run downloads it again.
|
||||
|
||||
@@ -55,6 +55,24 @@ inline constexpr int32_t kNotResident = -1;
|
||||
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);
|
||||
|
||||
/// #477 (--expert-profile-save): what the adaptive tier learned, as a profile ranking EVERY (layer, expert) pair -
|
||||
/// the experts resident in VRAM now first (where the swaps left the cache), then the rest; within each, by `heat`
|
||||
/// (the routing the adaptive tier counted since the start, descending), then by `prior` (the profile the engine
|
||||
/// started from: an expert this run never routed keeps its old place), then by index. A start from it begins
|
||||
/// where this one ended; one with more slots adds the hottest of the rest, one with fewer keeps the hottest.
|
||||
/// `resident` and `heat`: n_layers x n_expert entries (`heat` may be empty: no counts, the prior decides).
|
||||
std::vector<std::pair<int32_t, int32_t>> rank_learned_profile(int64_t n_layers, int64_t n_expert,
|
||||
const std::vector<uint8_t>& resident,
|
||||
const std::vector<double>& heat,
|
||||
const std::vector<std::pair<int32_t, int32_t>>& prior);
|
||||
|
||||
/// #477: writes `ranked` in tools/make_profile.py's format, byte for byte: `STRP`, `1, n_layers, n_expert,
|
||||
/// n_ranked, n_ranked` (uint32), the pairs (uint16 layer, uint16 expert), then the n_layers x n_expert int32 table
|
||||
/// of each pair's rank (-1: not ranked). Atomically: a `<path>.tmp` beside it, renamed over `path` once complete,
|
||||
/// so a reader (the next start) never sees half a file. False with `err` when it could not be written.
|
||||
bool write_expert_profile(const std::string& path, int64_t n_layers, int64_t n_expert,
|
||||
const std::vector<std::pair<int32_t, int32_t>>& ranked, std::string& err);
|
||||
|
||||
class ExpertCache {
|
||||
public:
|
||||
ExpertCache() = default;
|
||||
|
||||
@@ -421,6 +421,10 @@ public:
|
||||
/// CS-T, `budget_bytes` > 0 (`--resident-budget-gib`): only as many of those experts as fit `budget_bytes`, taken
|
||||
/// in `rank` order (the expert profile: the hottest after the GPU cache's), are copied; the rest stay on the
|
||||
/// mapped files (the SSD tier). No lend region then (a lent slot's expert is read from the files).
|
||||
/// #467: `budget_bytes` = `kResidentWhatFits` is that path sized by the RAM alone (available minus the headroom
|
||||
/// and the #403 margin) - the soft --resident-experts mode's second try when the whole complement does not fit;
|
||||
/// false when not even one expert fits. On Windows the mapped experts leave the working set before any reading.
|
||||
static constexpr uint64_t kResidentWhatFits = ~0ull;
|
||||
bool pin_cache_complement(
|
||||
const ExpertCache& cache, std::string& err, bool pin = true,
|
||||
const std::vector<std::pair<int32_t, int32_t>>& additional_gpu_pairs = {}, int64_t lend_from_slot = -1,
|
||||
|
||||
@@ -62,7 +62,18 @@
|
||||
#define cudaLaunchHostFunc hipLaunchHostFunc
|
||||
#define cudaMalloc hipMalloc
|
||||
#define cudaMallocHost(...) (::strata::hip_compat::malloc_host(__VA_ARGS__))
|
||||
#if defined(_WIN32)
|
||||
// Windows: hipMemGetInfo counts this process's own allocations only, so the desktop's and other programs' share of
|
||||
// the card is invisible to it, and VRAM sized by it overfills the card (WDDM then moves memory out to system RAM and
|
||||
// the GPU slows down). Its free figure is lowered by what Windows' video memory budget for this process withholds
|
||||
// (src/core/device.cu; STRATA_WDDM_BUDGET=0: hipMemGetInfo as it is).
|
||||
namespace strata::hip_compat {
|
||||
hipError_t mem_get_info(size_t* free_bytes, size_t* total_bytes);
|
||||
}
|
||||
#define cudaMemGetInfo ::strata::hip_compat::mem_get_info
|
||||
#else
|
||||
#define cudaMemGetInfo hipMemGetInfo
|
||||
#endif
|
||||
#define cudaMemcpy hipMemcpy
|
||||
#define cudaMemcpy2DAsync hipMemcpy2DAsync
|
||||
#define cudaMemcpyAsync hipMemcpyAsync
|
||||
|
||||
@@ -11,7 +11,8 @@
|
||||
// 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.
|
||||
// lowest index - the same selection as topk_kernel, over a quarter of the elements. On sm_90+
|
||||
// a call of up to 16 queries runs it on a cluster of 8 CTAs per query: the same ids.
|
||||
//
|
||||
// Queries carry their own step record (pos, n_kv, n_bid, width) as everywhere else in QSA.
|
||||
#pragma once
|
||||
@@ -41,6 +42,11 @@ void qsa_block_topk(const float* scores, const int32_t* steps, int64_t nq, int64
|
||||
///< the same contract as qsa_block_scores's). The register kernel is
|
||||
///< chosen by this, not by the capacity max_blocks (a long
|
||||
///< --max-context otherwise sends every short prompt to the slow one)
|
||||
/// The same ids on a thread-block cluster of 8 CTAs per query (sm_90+, CUDA; S19). qsa_block_topk takes it for calls
|
||||
/// of up to 16 queries unless STRATA_QSA_CLUSTER=0. False (nothing launched) where it cannot run: HIP, a card or a
|
||||
/// build below sm_90, or a capacity whose keys do not fit one cluster's shared memory. Capturable.
|
||||
bool qsa_block_topk_cluster(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);
|
||||
/// The original kernel (keys read from memory on every radix pass), for tests: the same ids.
|
||||
void qsa_block_topk_ref(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);
|
||||
|
||||
@@ -42,6 +42,11 @@ struct SamplerParams {
|
||||
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);
|
||||
|
||||
// The greedy pick (no penalties) on a thread-block cluster of 8 CTAs per row (sm_90+, CUDA; S19): the same token as
|
||||
// sample_tokens' one-block argmax, which takes it unless STRATA_ARGMAX_MULTI=0. False (nothing launched) where it
|
||||
// cannot run: HIP, a card or a build below sm_90. Device pointers; capturable.
|
||||
bool sample_greedy_cluster(const float* logits, int n_tokens, int n_vocab, int* out, void* stream);
|
||||
|
||||
// ---- COUPLED DRAFT SAMPLING (include/strata/core/coupled_draft.hpp, STRATA_SPEC_COUPLED=1). Device pointers
|
||||
// throughout; every per-request / per-round value comes from device memory, so the calls can be captured.
|
||||
//
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
// include/strata/prefill/moe_fused.hpp - #136 (prompt speed): the routed experts of a Q2_0 pack in the prompt path on
|
||||
// Strata's own int8 tensor-core kernels, opt-in with STRATA_PF_FUSED=1 (NVIDIA sm_80 and newer; sm_75 and HIP: MMQ).
|
||||
//
|
||||
// The MMQ path (moe_mmq.hpp) per MoE layer: a host sync to group the (token, k) pairs by expert, a gather of each
|
||||
// expert into a 16-expert group buffer, then per group the q8_1 activations of every routed slot, gate/up, SwiGLU,
|
||||
// the q8_1 of H and down - five passes over per-slot buffers. Here:
|
||||
// - the activations are rounded to int8 once per TOKEN (not per slot), with a scale per 32 values;
|
||||
// - the grouping (counts, offsets, the rows of each expert, a list of 64-row tiles) runs on the GPU from the
|
||||
// router's ids, so the host never waits for the routing;
|
||||
// - one launch per batch of up to kMaxBatch experts reads each expert where it already is (a resident cache slot
|
||||
// or a slot of the streamed ring: a pointer per expert, no gather) - gate/up with SwiGLU and the int8 rounding
|
||||
// of H in its epilogue, then down into the per-slot rows the existing combine reads.
|
||||
// The 2-bit codes are unpacked in registers and multiplied with mma.sync m16n8k32 (s8 x s8 -> s32); Q2_0's offset
|
||||
// (w = d * (q - 1)) is folded into the activation blocks' sums. The numbers differ from MMQ's (another int8
|
||||
// rounding of the activations and of H) by about as much as MMQ differs from FP32 (tests/cuda/prefill_fused_moe_test).
|
||||
#pragma once
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
|
||||
namespace strata::prefill::fused {
|
||||
|
||||
/// Routed rows (token, k pairs of one expert) per tile.
|
||||
constexpr int kTileRows = 64;
|
||||
/// Experts per launch: their blob pointers travel in the launch's parameters (1 KB).
|
||||
constexpr int kMaxBatch = 128;
|
||||
|
||||
/// This build has the kernels (a CUDA build; HIP and builds without MMQ do not).
|
||||
bool built();
|
||||
/// built(), and the current device is sm_80 or newer (mma.sync with s8 operands at m16n8k32).
|
||||
bool available();
|
||||
/// available(), unless STRATA_PF_FUSED=0 (read once): the Q2_0 pack's fused experts, on by default since 0.1.36.
|
||||
bool enabled();
|
||||
/// STRATA_PF_FUSED=1 given explicitly, and available(): the native IQ packs' fused kernels (opt-in).
|
||||
bool requested();
|
||||
|
||||
/// Bytes of `rows` activation rows of `cols` values (a multiple of 64) in the kernels' int8 form: per 64 values 64
|
||||
/// codes and {d, c} of each half (c = -(1.5 * 2^23 + the half's code sum), what the epilogue adds back).
|
||||
size_t act_bytes(int64_t rows, int64_t cols);
|
||||
/// Bytes of the grouping tables for `n` routed rows over `n_expert` experts.
|
||||
size_t group_bytes(int64_t n, int n_expert);
|
||||
|
||||
/// x [rows][cols] FP32 -> `xa` (act_bytes(rows, cols)).
|
||||
void quantize_act(const float* x, int64_t rows, int64_t cols, void* xa, void* stream);
|
||||
|
||||
/// The routing `ids` [n = tokens * k] -> per expert its rows (counts, offsets, the 64-row tile list, in `scratch`,
|
||||
/// group_bytes(n, n_expert)), `slot[i]` = the row of pair i (the row of the per-slot outputs the combine reads) and
|
||||
/// `src[row]` = its token (i / k). The order of the rows within an expert is not fixed; no output depends on it.
|
||||
void group(const int32_t* ids, int64_t n, int k, int n_expert, void* scratch, int32_t* slot, int32_t* src,
|
||||
void* stream);
|
||||
|
||||
/// Experts [e0, e1) of one layer (e1 - e0 <= kMaxBatch); blob[e - e0]: the Strata Q2_0 blob of expert e on the
|
||||
/// device (16-byte aligned), read only for an expert with rows.
|
||||
struct Batch {
|
||||
int e0 = 0, e1 = 0;
|
||||
const uint8_t* blob[kMaxBatch] = {};
|
||||
};
|
||||
/// The batch's products: gate/up from `xa` (quantize_act of the layer's input, [tokens][2560]) through the rows
|
||||
/// `src` of `group`, SwiGLU, H as int8 into `ha` (act_bytes(n, 640)), down into `dm` [n][2560] FP32 at the rows of
|
||||
/// `group`. `n`: the layer's routed rows (tokens * k), which bounds the launch grid.
|
||||
void experts(const Batch& b, int n_expert, int64_t n, const void* scratch, const void* xa, const int32_t* src, void* ha,
|
||||
float* dm, void* stream);
|
||||
|
||||
} // namespace strata::prefill::fused
|
||||
@@ -0,0 +1,41 @@
|
||||
// include/strata/prefill/moe_fused_iq.hpp - #136: the fused int8 prompt experts (moe_fused.hpp) for the native GGUF
|
||||
// packs (tools/iq_pack.py: an expert is its raw GGUF slices [gate rows | up rows | down rows], formats per layer in
|
||||
// expert_layout().fmt), opt-in with the same STRATA_PF_FUSED=1 (NVIDIA sm_80 and newer; sm_75 and HIP: MMQ).
|
||||
//
|
||||
// The grouping on the GPU (fused::group), the per-expert pointers (fused::Batch: a cache slot or a ring slot, no
|
||||
// gather) and the launch shape are the Q2_0 path's. What differs is the load stage: the i-quant blocks are decoded to
|
||||
// int8 in shared memory before the mma (codebook and sign lookups, as llama.cpp's MMQ load_tiles does), with the
|
||||
// block's scale per 32 (or per 16: IQ2_XS, IQ2_S, multiplied at m16n8k16) beside them. The activations and H are
|
||||
// int8 per 32 values in their natural order here (quantize_act_native), not the Q2_0 path's permuted one.
|
||||
//
|
||||
// Covered: gate/up IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, IQ3_S, IQ4_XS; down Q2_0, IQ4_NL - every layer of the IQ2_XS,
|
||||
// IQ3_XXS and IQ3_S packs but the IQ1_M ones (which MMQ does not cover either). Other layers keep MMQ.
|
||||
#pragma once
|
||||
|
||||
#include "strata/prefill/moe_fused.hpp"
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
|
||||
namespace strata::prefill::fused {
|
||||
|
||||
/// One layer's native expert geometry (cpu::NativeFmt's fields the kernels read; n_embd 2560, n_ff 640).
|
||||
struct NativeGeom {
|
||||
int gu_type = -1, d_type = -1; ///< ggml types
|
||||
size_t gu_row = 0, d_row = 0; ///< bytes per gate/up row and per down row
|
||||
size_t up_off = 0, down_off = 0; ///< inside the blob
|
||||
};
|
||||
|
||||
/// enabled() (STRATA_PF_FUSED=1 on sm_80+), and the kernels cover this gate/up and down pair on this device.
|
||||
bool native_supported(int gu_type, int d_type);
|
||||
|
||||
/// x [rows][cols] FP32 -> `xa` (act_bytes(rows, cols)): int8 per 32 values in natural order, the native kernels' form.
|
||||
void quantize_act_native(const float* x, int64_t rows, int64_t cols, void* xa, void* stream);
|
||||
|
||||
/// experts() for native blobs: b.blob[e - e0] is the native blob of expert e (gate at 0, up at g.up_off, down at
|
||||
/// g.down_off; 2-byte aligned). `xa` from quantize_act_native; H in `ha` (act_bytes(n, 640)); down into `dm` at the
|
||||
/// rows of `group`.
|
||||
void experts_native(const Batch& b, const NativeGeom& g, int n_expert, int64_t n, const void* scratch, const void* xa,
|
||||
const int32_t* src, void* ha, float* dm, void* stream);
|
||||
|
||||
} // namespace strata::prefill::fused
|
||||
+22
-3
@@ -143,10 +143,27 @@ def _late_system_to_user(messages: list[dict]) -> list[dict]:
|
||||
return [dict(m, role="user") if m.get("role") == "system" and i > 0 else m for i, m in enumerate(messages)]
|
||||
|
||||
|
||||
def _object_list(value, name: str) -> list[dict]:
|
||||
"""#460: a request's "messages" (or a message's "tool_calls") as a list of objects. Some clients send the array
|
||||
double-encoded, as a JSON string, which used to be iterated character by character and crashed on m.get: such a
|
||||
string is decoded. Anything that is still not a list of objects is a ValueError, which the server answers with
|
||||
a 400 naming the field. None (or no field) is an empty list, as a missing field always was."""
|
||||
if value is None:
|
||||
return []
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
value = json.loads(value)
|
||||
except ValueError:
|
||||
raise ValueError(f"{name} must be a list of objects (a string was sent that is not JSON)") from None
|
||||
if not isinstance(value, list) or not all(isinstance(m, dict) for m in value):
|
||||
raise ValueError(f"{name} must be a list of objects")
|
||||
return value
|
||||
|
||||
|
||||
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", []):
|
||||
for m in _object_list(req.get("messages"), "messages"):
|
||||
role = m.get("role")
|
||||
if role == "developer":
|
||||
role = "system"
|
||||
@@ -155,8 +172,10 @@ def openai_to_messages(req: dict) -> tuple[list[dict], list[dict] | None, dict]:
|
||||
out["reasoning_content"] = m["reasoning_content"]
|
||||
if m.get("tool_calls"):
|
||||
calls = []
|
||||
for c in m["tool_calls"]:
|
||||
for c in _object_list(m["tool_calls"], "tool_calls"):
|
||||
fn = c.get("function", c)
|
||||
if not isinstance(fn, dict):
|
||||
raise ValueError("tool_calls must be a list of objects (each with a \"function\" object)")
|
||||
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 {}
|
||||
@@ -186,7 +205,7 @@ def anthropic_to_messages(req: dict, think_unasked: bool = True) -> tuple[list[d
|
||||
system = req.get("system")
|
||||
if system:
|
||||
messages.append({"role": "system", "content": _text_of(system)})
|
||||
for m in req.get("messages", []):
|
||||
for m in _object_list(req.get("messages"), "messages"):
|
||||
content = m.get("content")
|
||||
if isinstance(content, str):
|
||||
messages.append({"role": m["role"], "content": content})
|
||||
|
||||
+264
-21
@@ -34,6 +34,7 @@ import re
|
||||
import select
|
||||
import signal
|
||||
import socket
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
@@ -65,6 +66,16 @@ CTX_SLACK = 8 # `strata --serve` rejects prompt + max_new + 8 > co
|
||||
# token (the Monitor showed five-digit numbers) and then undershoots for the first second of every answer.
|
||||
RATE_WINDOW_S = 2.0
|
||||
RATE_MIN_SPAN_S = 0.25 # younger than this there is no rate yet: the mean so far, with the span floored here
|
||||
# #481: a running request whose engine prints nothing (no T, PP or any other line) for this long has lost step with the
|
||||
# server (the engine's main thread waits, untimed, for its next command): the engine is ended and the request fails;
|
||||
# the next request starts it again. The config's "engine_silence_s" sets it (0: wait forever, as before).
|
||||
ENGINE_SILENCE_S = 300.0
|
||||
# ... except while a prompt is read: a PP line comes once per chunk (up to 32768 tokens with --prefill auto, issue
|
||||
# #282), and the slowest PCs read ~100 tok/s, so a first chunk can take minutes before the first line. Until the first
|
||||
# PP the wait adds the chunk's tokens at PP_FLOOR_TOK_S; after one, a chunk may take PP_SLACK x the last one's time.
|
||||
PP_CHUNK_MAX = 32768
|
||||
PP_FLOOR_TOK_S = 50.0
|
||||
PP_SLACK = 3.0
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ engines
|
||||
@@ -112,6 +123,12 @@ class ModelBusy(RuntimeError):
|
||||
"""Explicit model controls must not interrupt active or queued requests."""
|
||||
|
||||
|
||||
class EngineSilent(EngineDied):
|
||||
"""#481: the engine said nothing for too long during a request (or never acknowledged a STOP): the two sides lost
|
||||
step - the engine waiting for its next command, the server for this request's end - and the server ended it. An
|
||||
EngineDied, so the request ends with an error and the next one starts the engine again."""
|
||||
|
||||
|
||||
class EngineStuck(RuntimeError):
|
||||
"""The engine process did not end after QUIT, terminate and kill: the server keeps it (and says so) rather than
|
||||
reporting its GPU and RAM as given back."""
|
||||
@@ -123,12 +140,60 @@ class GpuBusy(RuntimeError):
|
||||
|
||||
|
||||
ENGINE_REQUEST = re.compile(
|
||||
r"prompt (?P<prompt>\d+) tokens = (?P<reused>\d+) reused \+ \d+ read in (?P<read>[\d.]+) ms \((?P<pp>[\d.]+) tok/s\), "
|
||||
# "+ 12288 of 98179 read": a request cancelled while its prompt was read (#471)
|
||||
r"prompt (?P<prompt>\d+) tokens = (?P<reused>\d+) reused \+ \d+(?: of \d+)? read in (?P<read>[\d.]+) ms "
|
||||
r"\((?P<pp>[\d.]+) tok/s\), "
|
||||
r"(?P<gen>\d+) generated in (?P<gen_ms>[\d.]+) ms \((?P<tg>[\d.]+) tok/s\)")
|
||||
|
||||
|
||||
_echoing: set[str] = set() # the logs echo_requests already follows (restart() runs StrataEngine.__init__ again)
|
||||
|
||||
DRAFT_HEAD_FAIL = "the draft head does not fit"
|
||||
DRAFT_HEAD_HINT = ("a smaller draft vocabulary needs less VRAM: --draft-vocab cyrillic (English, code and the Cyrillic "
|
||||
"script) or --draft-vocab en (English and code, ~215 MiB less than the default). Start once with "
|
||||
"it - START-HERE.bat --draft-vocab en (Windows) or ./setup.sh --draft-vocab en - and the model "
|
||||
"keeps it; or a smaller --context in setup.")
|
||||
|
||||
|
||||
def start_failure_hint(log: str | None, offset: int) -> str:
|
||||
"""#474: what to change when the engine stopped at the start because the MTP draft head did not fit the VRAM
|
||||
left: the engine's own `strata mtp:` lines after that failure (0.1.36+: what it needs, what is free, the smaller
|
||||
subsets), else the same advice in words for an older engine. "" for any other failure: the log says why."""
|
||||
if not log:
|
||||
return ""
|
||||
try:
|
||||
with open(log, "rb") as f:
|
||||
f.seek(offset)
|
||||
text = f.read().decode("utf-8", "replace")
|
||||
except (OSError, ValueError):
|
||||
return ""
|
||||
if DRAFT_HEAD_FAIL not in text:
|
||||
return ""
|
||||
said = [x.strip()[len("strata mtp: "):] for x in text.splitlines()
|
||||
if x.strip().startswith("strata mtp: ") and ("draft head over" in x or "hint:" in x)]
|
||||
return ". mtp: " + DRAFT_HEAD_FAIL + ". " + (" ".join(said) if said else "Hint: " + DRAFT_HEAD_HINT)
|
||||
|
||||
|
||||
def start_log_tail(log: str | None, offset: int, n: int = 20) -> str:
|
||||
"""#496: the last n lines this start wrote to the engine log, for the error when the engine ended before READY -
|
||||
whatever the failure, the engine's own reason is in them (people posted the traceback without the log). "" when
|
||||
there is no log or nothing in it from this start."""
|
||||
if not log:
|
||||
return ""
|
||||
try:
|
||||
with open(log, "rb") as f:
|
||||
f.seek(0, 2)
|
||||
start = max(offset, f.tell() - 64 * 1024) # enough for 20 lines, never an earlier start's
|
||||
f.seek(start)
|
||||
text = f.read().decode("utf-8", "replace")
|
||||
except (OSError, ValueError):
|
||||
return ""
|
||||
lines = text.splitlines()[1 if start > offset else 0:] # not a line cut in half
|
||||
lines = [x.rstrip() for x in lines if x.strip()][-n:]
|
||||
if not lines:
|
||||
return ""
|
||||
return "\nthe engine log's last lines:\n" + "\n".join(" " + x for x in lines)
|
||||
|
||||
|
||||
def echo_requests(log_path: str, offset: int) -> None:
|
||||
"""STRATA_REQUEST_LINES=1: one stdout line per finished request, from the engine's own summary in its log.
|
||||
@@ -153,6 +218,27 @@ def echo_requests(log_path: str, offset: int) -> None:
|
||||
m["tg"]), flush=True)
|
||||
|
||||
|
||||
def experts_loading_words(args: list, size: str) -> str:
|
||||
"""#505: what the start does with the experts, by the engine's flags (generate.cpp's option parsing): a RAM budget
|
||||
copies the hottest N GiB into RAM (--resident-budget-gib), the resident low-RAM mode the ones the GPU does not hold
|
||||
(--resident-experts), plain --mmap-experts reads them from the model files through the OS file cache (nothing is
|
||||
loaded into RAM up front); otherwise all of them go into RAM."""
|
||||
if "--resident-budget-gib" in args:
|
||||
try:
|
||||
n = f"up to {float(args[args.index('--resident-budget-gib') + 1]):g} GiB"
|
||||
except (IndexError, ValueError):
|
||||
n = "a RAM budget"
|
||||
return (f"loading the most-used experts into RAM ({n}; the rest are read from the model files as needed) "
|
||||
"and locking part of them for the GPU.")
|
||||
if "--resident-experts" in args:
|
||||
return (f"loading the experts the GPU does not hold into RAM (of {size}) and locking part of them for the "
|
||||
"GPU.")
|
||||
if "--mmap-experts" in args:
|
||||
return (f"mapping the experts from the model files ({size}, --mmap-experts): they are not loaded into RAM - "
|
||||
"the OS file cache reads them as the GPU's expert cache fills and as requests need them.")
|
||||
return f"loading the experts into RAM ({size}) and locking part of them for the GPU."
|
||||
|
||||
|
||||
def narrate_start(log_path: str, offset: int, args: list, done: threading.Event, heartbeat=20.0) -> None:
|
||||
"""While the engine starts, say in the server window what it is doing, from its log: the start reads tens of GB
|
||||
into RAM and locks part of it for the GPU, and on many PCs everything is slow or frozen for a minute or more -
|
||||
@@ -164,6 +250,7 @@ def narrate_start(log_path: str, offset: int, args: list, done: threading.Event,
|
||||
except (OSError, IndexError):
|
||||
pass
|
||||
size = f"about {gb:.0f} GB" if gb >= 1 else "tens of GB"
|
||||
loading = experts_loading_words(args, size)
|
||||
t0 = last = time.time()
|
||||
said = set()
|
||||
|
||||
@@ -187,8 +274,8 @@ def narrate_start(log_path: str, offset: int, args: list, done: threading.Event,
|
||||
cut = chunk.rfind(b"\n") + 1
|
||||
pos += cut
|
||||
for line in chunk[:cut].decode("utf-8", "replace").splitlines():
|
||||
if "PLE on" in line or "expert arena:" in line:
|
||||
say("arena", f"[strata] loading the experts into RAM ({size}) and locking part of them for the GPU.\n"
|
||||
if "PLE on" in line or "expert arena:" in line or "experts via mmap" in line: # #505: mapped
|
||||
say("arena", f"[strata] {loading}\n"
|
||||
" YOUR PC CAN BE SLOW OR STOP RESPONDING FOR 1-3 MINUTES NOW - this is normal.\n"
|
||||
" Please wait and don't close this window; the browser opens when it is ready.")
|
||||
elif " loaded " in line and "GiB at" in line:
|
||||
@@ -212,6 +299,8 @@ class StrataEngine:
|
||||
(`temperature=F top_p=F top_k=N seed=N`, the engine's own spelling). An absent temperature keeps the
|
||||
engine's default, which is greedy; `temperature=0` means the same thing, so it is not forwarded.
|
||||
"""
|
||||
silence_s = ENGINE_SILENCE_S # #481: main() sets the config's engine_silence_s (survives restart())
|
||||
silent_note = None # #481: why the server ended a silent engine (death_note says it)
|
||||
|
||||
def __init__(self, exe: str, args: list[str], cwd: str | None = None, log: str | None = None,
|
||||
env: dict | None = None, lazy: bool = False):
|
||||
@@ -227,7 +316,8 @@ class StrataEngine:
|
||||
self.info = {} # INFO key=value facts (engine 0.1.8+): kv, expert slots, ... (Monitor tab)
|
||||
self.prefill_tok_s_mean = None
|
||||
self.progress = None # (read, total) prompt tokens while a prompt is read, from PP lines
|
||||
try: # a ready-made engine's BUILD.json says its version
|
||||
self.silent_note = None
|
||||
try: # a ready-made engine's BUILD.json says its version
|
||||
self.info["version"] = json.loads((Path(exe).parent / "BUILD.json").read_text()).get("version")
|
||||
except (OSError, ValueError):
|
||||
self.info["version"] = None
|
||||
@@ -236,6 +326,7 @@ class StrataEngine:
|
||||
self.unloaded = False # `ended` stays True until READY (below): not alive while starting (#344)
|
||||
self.log = open(log, "a", encoding="utf-8") if log else subprocess.DEVNULL
|
||||
loading = threading.Event() # set once READY: the narrator below stops
|
||||
log_start = os.path.getsize(log) if log else 0 # where this start's lines begin (start_failure_hint)
|
||||
if log:
|
||||
threading.Thread(target=narrate_start, args=(log, os.path.getsize(log), args, loading),
|
||||
daemon=True).start()
|
||||
@@ -259,7 +350,16 @@ class StrataEngine:
|
||||
break
|
||||
loading.set()
|
||||
if self.max_context <= 0:
|
||||
raise RuntimeError("the engine exited before it was ready" + (f" (see {log})" if log else ""))
|
||||
try: # its pipes and our handle on its log (the log stays)
|
||||
self.proc.wait(timeout=5)
|
||||
self.proc.stdin.close()
|
||||
self.proc.stdout.close()
|
||||
if log:
|
||||
self.log.close()
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
raise RuntimeError("the engine exited before it was ready" + (f" (see {log})" if log else "") +
|
||||
start_failure_hint(log, log_start) + start_log_tail(log, log_start))
|
||||
# (from PR #41, midhatn) a locally built engine can sit next to another release's BUILD.json: engines that
|
||||
# report their own version (INFO engine=, 0.1.8+) win, the manifest stays the fallback for older ones
|
||||
if self.info.get("engine"):
|
||||
@@ -280,6 +380,8 @@ class StrataEngine:
|
||||
|
||||
def death_note(self) -> str:
|
||||
"""Why the engine most likely ended, from the end of its log: its own watchdog (issue #29), else RAM."""
|
||||
if getattr(self, "silent_note", None): # #481: the server ended it, not the OS or the engine itself
|
||||
return self.silent_note
|
||||
tail = ""
|
||||
try:
|
||||
with open(self.log_path, "rb") as f:
|
||||
@@ -338,6 +440,8 @@ class StrataEngine:
|
||||
self.last.update(hits=int(f[9]), lookups=int(f[10]))
|
||||
if len(f) >= 14: # the expert tiers (engine 0.1.31+): RAM / file blobs, file MB
|
||||
self.last.update(ram_blobs=int(f[11]), file_blobs=int(f[12]), file_mb=float(f[13]))
|
||||
if len(f) >= 15: # #471 (engine 0.1.36+): the prompt tokens actually read
|
||||
self.last.update(prompt_read=int(f[14]))
|
||||
|
||||
@staticmethod
|
||||
def sampling_keys(sampling: dict) -> str:
|
||||
@@ -408,19 +512,34 @@ class StrataEngine:
|
||||
except OSError: # the pipe is gone: the engine died (not the client)
|
||||
raise EngineDied(f"the engine stopped unexpectedly (exit code {self.exit_code()})") from None
|
||||
done = False
|
||||
# #481: how long the engine may stay silent from here. Until the first line: the request's first prompt
|
||||
# chunk at the slowest prompt reading on top of silence_s; a PP line resets it to its own chunk's time.
|
||||
silence = float(self.silence_s or 0)
|
||||
allow = silence + min(len(ids), PP_CHUNK_MAX) / PP_FLOOR_TOK_S if silence > 0 else 0.0
|
||||
heard, read_to = time.monotonic(), 0
|
||||
try:
|
||||
while True:
|
||||
wait = 10.0
|
||||
if allow > 0:
|
||||
wait = min(wait, allow - (time.monotonic() - heard))
|
||||
if wait <= 0:
|
||||
done = True # no STOP and no drain: nothing is listening
|
||||
raise self._silent(f"the engine said nothing for {time.monotonic() - heard:.0f} s during "
|
||||
"the request")
|
||||
try:
|
||||
line = self.lines.get(timeout=10)
|
||||
line = self.lines.get(timeout=wait)
|
||||
except queue.Empty:
|
||||
if cancel.is_set():
|
||||
return
|
||||
yield None
|
||||
if wait >= 10.0:
|
||||
yield None # the 10 s heartbeat (the deadline's short waits are not)
|
||||
continue
|
||||
if line is None:
|
||||
done = True
|
||||
raise EngineDied(f"the engine stopped unexpectedly (exit code {self.exit_code()})")
|
||||
heard = time.monotonic() # any line is output: T, PP, RESUME, INFO ...
|
||||
if line.startswith("T "):
|
||||
allow = silence
|
||||
if cancel.is_set():
|
||||
return
|
||||
yield int(line[2:])
|
||||
@@ -429,9 +548,18 @@ class StrataEngine:
|
||||
if len(f) >= 3 and f[1].isdigit() and f[2].isdigit():
|
||||
self.progress = (int(f[1]), int(f[2])) # prompt progress, one per chunk: also a heartbeat (the
|
||||
self.prefill_tok_s_mean = float(f[4]) if len(f) >= 5 else None
|
||||
rate, chunk = self.prefill_tok_s_mean or 0.0, int(f[1]) - read_to
|
||||
read_to = int(f[1])
|
||||
if silence > 0 and rate > 0 and chunk > 0: # #481: the next chunk, as long as this one
|
||||
allow = max(silence, PP_SLACK * chunk / rate)
|
||||
if cancel.is_set(): # lines reset the 10 s wait, so without this a long prompt
|
||||
return # would send no keep-alives at all)
|
||||
yield None
|
||||
elif line.startswith("RESUME "): # the reused tokens: the first chunk starts after them
|
||||
try:
|
||||
read_to = int(line.split()[1])
|
||||
except (IndexError, ValueError):
|
||||
pass
|
||||
elif line.startswith("DONE"):
|
||||
self._parse_done(line)
|
||||
done = True
|
||||
@@ -447,13 +575,41 @@ class StrataEngine:
|
||||
self.proc.stdin.flush()
|
||||
except OSError:
|
||||
pass
|
||||
# #481: never an untimed wait here - it holds the request FIFO, and an engine that lost step never
|
||||
# answers. An engine that honours STOP gets the current allowance in all (a STOP during a prompt
|
||||
# chunk is seen after it); an older one runs on to max_new, so each line only has to come in time.
|
||||
heard = time.monotonic()
|
||||
while True:
|
||||
line = self.lines.get()
|
||||
left = allow - (time.monotonic() - heard) if allow > 0 else None
|
||||
try:
|
||||
if left is not None and left <= 0:
|
||||
raise queue.Empty
|
||||
line = self.lines.get(timeout=left)
|
||||
except queue.Empty:
|
||||
raise self._silent("the engine did not finish the request after it was stopped (STOP) "
|
||||
f"within {allow:.0f} s") from None
|
||||
if line is None or line.startswith("ERR"):
|
||||
break
|
||||
if line.startswith("DONE"):
|
||||
self._parse_done(line)
|
||||
break
|
||||
if not self.can_stop:
|
||||
heard = time.monotonic()
|
||||
|
||||
def _silent(self, what: str) -> EngineSilent:
|
||||
"""#481: end an engine that lost step with the server (its main thread waits for a command the server never
|
||||
sends), so the next request starts it again: killed now, its GPU and RAM go back with the process."""
|
||||
self.silent_note = ("The engine and the server lost step (issue #481; a very slow PC can raise "
|
||||
"\"engine_silence_s\" in the config, 0 = wait forever). If it happens again, please add "
|
||||
"the end of the engine log to github.com/Niko1221/Strata/issues/481.")
|
||||
self.ended = True # not alive from now: the next request restarts it
|
||||
proc = self.proc
|
||||
try:
|
||||
proc.kill()
|
||||
proc.wait(timeout=20) # restart() -> close() handles one that is still exiting
|
||||
except (OSError, AttributeError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
return EngineSilent(f"{what}; the server ended the engine")
|
||||
|
||||
def close(self):
|
||||
"""End the engine process: QUIT first (the engine frees its memory itself - unpinning tens of GB can take
|
||||
@@ -621,6 +777,17 @@ def gpu_list(cfg: dict) -> list[int]:
|
||||
return [int(str(x).strip()) for x in items if str(x).strip() != ""]
|
||||
|
||||
|
||||
def engine_silence_s(cfg: dict) -> float:
|
||||
"""#481: the config's "engine_silence_s" - seconds an engine may print nothing during a request before the server
|
||||
ends it (default ENGINE_SILENCE_S; 0 = wait forever). ValueError for anything but a number >= 0."""
|
||||
v = cfg.get("engine_silence_s")
|
||||
if v is None:
|
||||
return ENGINE_SILENCE_S
|
||||
if isinstance(v, bool) or not isinstance(v, (int, float)) or v < 0:
|
||||
raise ValueError(f'"engine_silence_s" must be a number of seconds >= 0 (0 = no limit), not {v!r}')
|
||||
return float(v)
|
||||
|
||||
|
||||
def engine_args(cfg: dict) -> list[str]:
|
||||
"""The engine's arguments: the config's, and with several GPUs the layer split across them ("layer_split" in the
|
||||
config: "auto" by default, or the first layer of each later GPU's share, e.g. "18" or "16,32")."""
|
||||
@@ -630,6 +797,44 @@ def engine_args(cfg: dict) -> list[str]:
|
||||
# opt-in: an auto split runs on the first card alone when it holds every profiled expert and the KV
|
||||
if len(gpu_list(cfg)) > 1 and cfg.get("split_skip_if_fits") and "--split-skip-if-fits" not in args:
|
||||
args.append("--split-skip-if-fits")
|
||||
return learned_profile_args(cfg, args)
|
||||
|
||||
|
||||
def profile_shape(path: str) -> tuple[int, int] | None:
|
||||
"""An expert profile's (layers, experts per layer), from its header (tools/make_profile.py's format), or None
|
||||
when the file is missing, is not one or is shorter than the pairs its header promises."""
|
||||
try:
|
||||
with open(path, "rb") as f:
|
||||
head = f.read(24)
|
||||
size = os.fstat(f.fileno()).st_size
|
||||
except OSError:
|
||||
return None
|
||||
if len(head) < 24 or head[:4] != b"STRP":
|
||||
return None
|
||||
_, nl, ne, _, n = struct.unpack("<5I", head[4:])
|
||||
return (nl, ne) if size >= 24 + 4 * n else None
|
||||
|
||||
|
||||
def learned_profile_args(cfg: dict, args: list[str]) -> list[str]:
|
||||
"""#477 (opt-in): "expert_profile_save": "<path>" in the config has the engine (0.1.36+) save what its adaptive
|
||||
tier learned there - on QUIT and every "expert_profile_save_every" minutes (10 by default, 0 = at QUIT only) -
|
||||
and the next start begins from it instead of the config's --expert-profile, when it is a profile of the same
|
||||
model (its header's layers and experts match); otherwise from the config's own, as before. A relative path is
|
||||
the engine's (the config's "cwd"). Without the key, the arguments are the config's, unchanged."""
|
||||
save = cfg.get("expert_profile_save")
|
||||
if not isinstance(save, str) or not save.strip() or "--expert-profile-save" in args:
|
||||
return args
|
||||
args = args + ["--expert-profile-save", save]
|
||||
every = cfg.get("expert_profile_save_every")
|
||||
if isinstance(every, (int, float)) and not isinstance(every, bool) and every >= 0:
|
||||
args += ["--expert-profile-save-every", str(every)]
|
||||
if "--expert-profile" in args[:-1]:
|
||||
i = args.index("--expert-profile") + 1
|
||||
here = cfg.get("cwd") or "."
|
||||
learned = profile_shape(save if os.path.isabs(save) else os.path.join(here, save))
|
||||
base = profile_shape(args[i] if os.path.isabs(args[i]) else os.path.join(here, args[i]))
|
||||
if learned is not None and learned == base:
|
||||
args[i] = save
|
||||
return args
|
||||
|
||||
|
||||
@@ -765,7 +970,8 @@ class Service:
|
||||
self.history = collections.deque(maxlen=500) # the last finished requests, newest last (GET /metrics)
|
||||
# since the server started (the Monitor's totals, issue #35)
|
||||
self.totals = {"since": time.time(), "requests": 0, "prompt_tokens": 0, "reused": 0, "output_tokens": 0,
|
||||
"prompt_ms": 0.0, "decode_ms": 0.0}
|
||||
"prompt_ms": 0.0, "decode_ms": 0.0,
|
||||
"drafts_offered": 0, "drafts_accepted": 0} # #457: the MTP drafts, summed where reported
|
||||
self.last_timings = None # the last finished request's, llama.cpp's names (/v1/status)
|
||||
self.last_request_at = None # when a request last started or finished
|
||||
self.started_at = time.time()
|
||||
@@ -878,6 +1084,13 @@ class Service:
|
||||
self.engine.restart()
|
||||
print("[strata] the engine is running again", flush=True)
|
||||
|
||||
def _say_died(self, e: Exception) -> None:
|
||||
"""The server window's line for an engine that died (or was ended, #481) in the middle of a request."""
|
||||
note = self.engine.death_note() if hasattr(self.engine, "death_note") else ""
|
||||
log = getattr(self.engine, "log_path", None)
|
||||
print(f"[strata] {e}. {note} The next request starts the engine again."
|
||||
f"{' Its log: ' + log if log else ''}", flush=True)
|
||||
|
||||
def load(self):
|
||||
"""POST /load and every generation request: start the engine now if it is unloaded (raises GpuBusy)."""
|
||||
# a request is on its way: the idle thread must not unload between this and the request's own start
|
||||
@@ -1248,7 +1461,7 @@ class Service:
|
||||
while True:
|
||||
gen = self.engine.generate(prompt, max_new - n, sampling, cancel, embeddings=emb) if emb \
|
||||
else self.engine.generate(prompt, max_new - n, sampling, cancel)
|
||||
seg, wrap = [], False # this pass's tokens; the budget is reached
|
||||
seg, wrap, leaving = [], False, False # this pass's tokens; the budget is reached; closed
|
||||
try:
|
||||
for t in gen:
|
||||
if t is None: # heartbeat while the engine is quiet
|
||||
@@ -1278,19 +1491,25 @@ class Service:
|
||||
break
|
||||
except EngineDied as e:
|
||||
finish = "error"
|
||||
note = self.engine.death_note() if hasattr(self.engine, "death_note") else ""
|
||||
log = getattr(self.engine, "log_path", None)
|
||||
print(f"[strata] {e}. {note} The next request starts the engine again."
|
||||
f"{' Its log: ' + log if log else ''}", flush=True)
|
||||
self._say_died(e)
|
||||
raise
|
||||
except ValueError as e: # the engine's ERR line (it may have ended after it)
|
||||
finish = "error"
|
||||
print(f"[strata] the engine reported an error: {e}", flush=True)
|
||||
raise
|
||||
except GeneratorExit: # the client went away: an engine that never acknowledges
|
||||
leaving = True # the STOP below is ended, but no error replaces this
|
||||
raise
|
||||
finally:
|
||||
gen.close() # STOP+drain to THIS request's DONE while still holding the
|
||||
# fifo, so a stop-token break can't leave the shared engine
|
||||
# queue mid-drain for the next request to read as its own DONE
|
||||
try:
|
||||
gen.close() # STOP+drain to THIS request's DONE while still holding the
|
||||
# fifo, so a stop-token break can't leave the shared engine
|
||||
# queue mid-drain for the next request to read as its own DONE
|
||||
except EngineSilent as e: # #481: the STOP was never acknowledged: the engine is ended
|
||||
finish = "error"
|
||||
self._say_died(e)
|
||||
if not leaving and not cancel.is_set():
|
||||
raise
|
||||
if not wrap or cancel.is_set():
|
||||
break
|
||||
# #123: the thinking reached reasoning_budget_tokens. Close it the way the model would (a
|
||||
@@ -1328,27 +1547,35 @@ class Service:
|
||||
cvec = (getattr(self.engine, "info", {}) or {}).get("cvec", 0)
|
||||
loaded = str(cvec) not in ("0", "", "None")
|
||||
hit_rate = round(last["hits"] / last["lookups"], 3) if last.get("lookups") else None
|
||||
seen = prompt_tokens_seen(len(ids), last) # #471: < len(ids) when cancelled mid-read
|
||||
self.history.append({
|
||||
"projection": (sampling or {}).get("experimental_speed_projection") is not False
|
||||
if loaded else None,
|
||||
"time": started, "duration_s": round(time.time() - started, 1), "finish": finish,
|
||||
"prompt_tokens": len(ids), "reused": last.get("reused"), "output_tokens": n,
|
||||
"prompt_tokens": seen, "reused": last.get("reused"), "output_tokens": n,
|
||||
# the request's whole prompt, and the tokens read of it (None: an older engine)
|
||||
"prompt_total": len(ids), "prompt_read": last.get("prompt_read"),
|
||||
"engine_generated": last.get("generated"),
|
||||
"prompt_ms": last.get("prompt_ms"), "decode_ms": last.get("decode_ms"),
|
||||
"decode_tok_s": round(last["generated"] / (last["decode_ms"] / 1000), 1)
|
||||
if n and last.get("generated") and last.get("decode_ms") else None,
|
||||
"hit_rate": hit_rate, "ram_blobs": last.get("ram_blobs"),
|
||||
"file_blobs": last.get("file_blobs"), "file_mb": last.get("file_mb")})
|
||||
"file_blobs": last.get("file_blobs"), "file_mb": last.get("file_mb"),
|
||||
# #457: the speculative drafts from the DONE line (None: the engine did not say)
|
||||
"drafts_offered": last.get("drafts_offered"),
|
||||
"drafts_accepted": last.get("drafts_accepted")})
|
||||
t = self.totals
|
||||
t["requests"] += 1
|
||||
t["prompt_tokens"] += len(ids)
|
||||
t["prompt_tokens"] += seen
|
||||
t["reused"] += last.get("reused") or 0
|
||||
t["output_tokens"] += n
|
||||
t["prompt_ms"] += last.get("prompt_ms") or 0.0
|
||||
t["decode_ms"] += last.get("decode_ms") or 0.0
|
||||
t["drafts_offered"] += last.get("drafts_offered") or 0
|
||||
t["drafts_accepted"] += last.get("drafts_accepted") or 0
|
||||
fresh = getattr(self.engine, "last", None)
|
||||
if fresh is not None and fresh is not before: # the engine's clock for THIS request
|
||||
timings = request_timings(len(ids), n, last)
|
||||
timings = request_timings(seen, n, last)
|
||||
self.last_timings = dict(timings, at=int(time.time())) if timings else None
|
||||
self.last_request_at = time.time()
|
||||
now = time.time()
|
||||
@@ -1372,6 +1599,17 @@ class Service:
|
||||
"timings": timings}
|
||||
|
||||
|
||||
def prompt_tokens_seen(prompt_tokens: int, last: dict) -> int:
|
||||
"""#471: the prompt tokens a request got through - all of them, unless the engine's DONE line says a cancel stopped
|
||||
its prompt read part-way (then the reused ones plus those read). /metrics' history and totals count these, so a
|
||||
cancelled read is neither recorded as the whole prompt nor given a rate from tokens it never read. An engine
|
||||
before 0.1.36 does not say (no `prompt_read`): the whole prompt, as before."""
|
||||
read = last.get("prompt_read")
|
||||
if read is None or last.get("finish") != "cancel":
|
||||
return prompt_tokens
|
||||
return min(prompt_tokens, int(last.get("reused") or 0) + int(read))
|
||||
|
||||
|
||||
def request_timings(prompt_tokens: int, generated: int, last: dict) -> dict | None:
|
||||
"""One request's `timings` in llama.cpp's names (what its clients show as speed), from the engine's own clock
|
||||
(StrataEngine.last): prompt_n is what was read, cache_n what the conversation cache already held. None when the
|
||||
@@ -2487,7 +2725,12 @@ def main() -> int:
|
||||
# a relative "exe" is the config's cwd's: Windows' CreateProcess resolves "engine/strata.exe" against nothing
|
||||
# it is told about (WinError 2), so it is made absolute here
|
||||
exe = cfg["exe"] if os.path.isabs(cfg["exe"]) else os.path.abspath(os.path.join(cfg.get("cwd") or ".", cfg["exe"]))
|
||||
try:
|
||||
silence = engine_silence_s(cfg) # #481: checked before the (minutes-long) start
|
||||
except ValueError as e:
|
||||
raise SystemExit(f"[strata] config {e}")
|
||||
engine = StrataEngine(exe, engine_args(cfg), cwd=cfg.get("cwd"), log=cfg.get("log"), env=env, lazy=lazy)
|
||||
engine.silence_s = silence # an attribute of its own: restart() keeps it
|
||||
warn_tight_ram(engine.info.get("arena_mib"))
|
||||
else:
|
||||
engine, vision, sampling_defaults = MockEngine(tok, a.script or [
|
||||
|
||||
+459
-1
@@ -18,7 +18,9 @@ from unittest import mock
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
from serve.frontend import ChatTemplate # noqa: E402
|
||||
from serve.server import CTX_SLACK, ByteTokenizer, EngineDied, GpuBusy, MockEngine, Service, StrataEngine, request_timings, serve # noqa: E402
|
||||
from serve.server import (CTX_SLACK, ByteTokenizer, EngineDied, GpuBusy, MockEngine, Service, StrataEngine, # noqa: E402
|
||||
engine_args, prompt_tokens_seen, request_timings, serve, start_failure_hint)
|
||||
from types import SimpleNamespace # noqa: E402
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
CTX = 4096
|
||||
@@ -95,6 +97,13 @@ class MaxTokens(unittest.TestCase):
|
||||
m = ENGINE_REQUEST.search(line)
|
||||
self.assertIsNotNone(m)
|
||||
self.assertEqual((m["prompt"], m["reused"], m["gen"], m["tg"]), ("1200", "1000", "30", "100.0"))
|
||||
# #471: a request cancelled while its prompt was read says how far it got
|
||||
m = ENGINE_REQUEST.search("strata serve: prompt 98179 tokens = 0 reused + 12288 of 98179 read in 17565 ms "
|
||||
"(699.6 tok/s), 0 generated in 0 ms (0.0 tok/s), drafts accepted 0 of 0, "
|
||||
"0 checkpoints (cancelled)")
|
||||
self.assertIsNotNone(m)
|
||||
self.assertEqual((m["prompt"], m["reused"], m["read"], m["pp"], m["gen"]),
|
||||
("98179", "0", "17565", "699.6", "0"))
|
||||
|
||||
def test_unset_budget_is_the_rest_of_the_context(self):
|
||||
cases = {"openai": [{"max_tokens": -1}, {"max_tokens": 0}, {}, {"max_tokens": None},
|
||||
@@ -474,6 +483,18 @@ class ClientShapes(unittest.TestCase):
|
||||
"messages": [{"role": "user", "content": "hi"}]})
|
||||
self.assertEqual(status, 200, b) # the server goes on
|
||||
|
||||
def test_messages_as_a_json_string_over_http(self):
|
||||
# #460: a double-encoded "messages" is answered; one that is not a list of objects is a 400, not a 500
|
||||
encoded = json.dumps([{"role": "user", "content": "1+1?"}])
|
||||
for path in ("/v1/chat/completions", "/v1/messages"):
|
||||
with self.subTest(path=path):
|
||||
status, b = self.post(path, {"model": "x", "max_tokens": 20, "messages": encoded})
|
||||
self.assertEqual(status, 200, b)
|
||||
self.assertIn("1+1?", self.prompt_text())
|
||||
status, b = self.post(path, {"model": "x", "max_tokens": 20, "messages": ["hi"]})
|
||||
self.assertEqual(status, 400, b)
|
||||
self.assertIn("messages must be a list of objects", b["error"]["message"])
|
||||
|
||||
def test_vision_temp_image_removed_when_the_pipe_fails(self):
|
||||
# #352: the temporary image goes even when the encoder's pipe raises
|
||||
from serve.server import Vision
|
||||
@@ -499,6 +520,32 @@ class ClientShapes(unittest.TestCase):
|
||||
msgs, _, _ = anthropic_to_messages({"system": "S", "messages": [{"role": "user", "content": "u"}]})
|
||||
self.assertEqual([m["role"] for m in msgs], ["system", "user"])
|
||||
|
||||
def test_messages_sent_as_a_json_string(self):
|
||||
# #460: a client that double-encodes "messages" (and "tool_calls") as a JSON string gets them decoded; what is
|
||||
# still not a list of objects is a ValueError (the server's 400), not an AttributeError on m.get
|
||||
from serve.frontend import anthropic_to_messages, openai_to_messages
|
||||
call = [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{\"x\": 1}"}}]
|
||||
listed = [{"role": "user", "content": "u"}, {"role": "assistant", "content": "", "tool_calls": call}]
|
||||
encoded = [listed[0], dict(listed[1], tool_calls=json.dumps(call))]
|
||||
want = openai_to_messages({"messages": listed})[0]
|
||||
self.assertEqual(want[1]["tool_calls"], [{"function": {"name": "f", "arguments": {"x": 1}}}])
|
||||
self.assertEqual(openai_to_messages({"messages": json.dumps(listed)})[0], want)
|
||||
self.assertEqual(openai_to_messages({"messages": json.dumps(encoded)})[0], want)
|
||||
anth = [{"role": "user", "content": "u"}]
|
||||
self.assertEqual(anthropic_to_messages({"messages": json.dumps(anth)})[0],
|
||||
anthropic_to_messages({"messages": anth})[0])
|
||||
self.assertEqual(openai_to_messages({})[0], []) # no field: nothing, as before
|
||||
self.assertEqual(openai_to_messages({"messages": None})[0], [])
|
||||
for bad in ("not json", "\"a string\"", json.dumps({"role": "user"}), ["hi"], [{"role": "user"}, 3], 5,
|
||||
{"role": "user", "content": "u"}):
|
||||
for fn in (openai_to_messages, anthropic_to_messages):
|
||||
with self.subTest(bad=bad, fn=fn.__name__):
|
||||
with self.assertRaisesRegex(ValueError, "messages must be a list of objects"):
|
||||
fn({"messages": bad})
|
||||
for calls in (["f"], "[1]", [{"function": "f"}], "{"):
|
||||
with self.subTest(calls=calls), self.assertRaisesRegex(ValueError, "tool_calls must be a list of objects"):
|
||||
openai_to_messages({"messages": [{"role": "assistant", "content": "", "tool_calls": calls}]})
|
||||
|
||||
|
||||
class SamplingKeys(unittest.TestCase):
|
||||
"""The GEN line's sampling keys: top_k 0 ("off") or wider than the engine's 64 get the widest list, 64 (they used
|
||||
@@ -709,6 +756,259 @@ class EngineDeath(unittest.TestCase):
|
||||
httpd.server_close()
|
||||
|
||||
|
||||
class DoneLineEngine(MockEngine):
|
||||
"""The mock engine whose `last` comes from a DONE line, parsed as StrataEngine parses it."""
|
||||
|
||||
def __init__(self, *a, done_lines=(), **kw):
|
||||
super().__init__(*a, **kw)
|
||||
self.done_lines = list(done_lines)
|
||||
|
||||
def generate(self, ids, max_new, sampling, cancel, embeddings=None):
|
||||
try:
|
||||
yield from super().generate(ids, max_new, sampling, cancel, embeddings)
|
||||
finally:
|
||||
StrataEngine._parse_done(self, self.done_lines.pop(0))
|
||||
|
||||
|
||||
class DraftCounts(unittest.TestCase):
|
||||
"""#457: GET /metrics gives each request's speculative draft counts (offered / accepted, from the engine's DONE
|
||||
line; None when the line has no such fields) and their running sums in the totals."""
|
||||
|
||||
def test_drafts_in_history_and_totals(self):
|
||||
tok = ByteTokenizer()
|
||||
engine = DoneLineEngine(tok, "</think>\n\nok", max_context=CTX, done_lines=[
|
||||
"DONE 4 20 40.0 30.0 stop 7 12 0", # 7 of 12 drafts accepted
|
||||
"DONE 4 20 40.0 30.0 stop", # an engine that reports no drafts
|
||||
"DONE 4 20 40.0 30.0 stop 3 5 0 9 10"])
|
||||
svc = Service(engine, tok, ChatTemplate(ROOT / "serve/chat_template.jinja"))
|
||||
self.assertEqual((svc.totals["drafts_offered"], svc.totals["drafts_accepted"]), (0, 0))
|
||||
httpd = serve(svc, port=0)
|
||||
base = f"http://127.0.0.1:{httpd.server_address[1]}"
|
||||
try:
|
||||
for _ in range(3):
|
||||
body = json.dumps({"model": "m", "max_tokens": 10, "messages": [{"role": "user", "content": "hi"}]})
|
||||
with urllib.request.urlopen(urllib.request.Request(base + "/v1/chat/completions", data=body.encode(),
|
||||
headers={"Content-Type": "application/json"}),
|
||||
timeout=30) as r:
|
||||
self.assertEqual(r.status, 200)
|
||||
with urllib.request.urlopen(base + "/metrics", timeout=10) as r:
|
||||
m = json.loads(r.read())
|
||||
finally:
|
||||
httpd.shutdown()
|
||||
httpd.server_close()
|
||||
rows = m["requests"] # newest first
|
||||
self.assertEqual([(r["drafts_offered"], r["drafts_accepted"]) for r in rows], [(5, 3), (None, None), (12, 7)])
|
||||
self.assertEqual((m["totals"]["drafts_offered"], m["totals"]["drafts_accepted"]), (17, 10))
|
||||
|
||||
|
||||
class LearnedProfile(unittest.TestCase):
|
||||
"""#477: "expert_profile_save" in the config: the engine saves its learned profile there, and the next start
|
||||
begins from it when it is a profile of the same model; without the key the arguments are unchanged."""
|
||||
|
||||
def write(self, path, nl=48, ne=512, n=4):
|
||||
sys.path.insert(0, str(ROOT / "tools"))
|
||||
import make_profile
|
||||
old = make_profile.N_LAYER
|
||||
make_profile.N_LAYER = nl
|
||||
try:
|
||||
make_profile.write_profile(path, [(i % nl, i // nl) for i in range(n)], n_expert=ne)
|
||||
finally:
|
||||
make_profile.N_LAYER = old
|
||||
|
||||
def test_without_the_key_nothing_changes(self):
|
||||
args = ["--native", "x", "--expert-profile", "data/expert-profile.bin", "--adapt-every", "4"]
|
||||
self.assertEqual(engine_args({"args": list(args)}), args)
|
||||
self.assertEqual(engine_args({"args": list(args), "expert_profile_save": ""}), args)
|
||||
|
||||
def test_save_and_start_from_it(self):
|
||||
d = Path(tempfile.mkdtemp())
|
||||
self.write(d / "base.bin")
|
||||
cfg = {"args": ["--expert-profile", "base.bin"], "cwd": str(d), "expert_profile_save": "learned.bin",
|
||||
"expert_profile_save_every": 5}
|
||||
# nothing saved yet: the config's profile, and the engine is told where to save
|
||||
self.assertEqual(engine_args(cfg), ["--expert-profile", "base.bin", "--expert-profile-save", "learned.bin",
|
||||
"--expert-profile-save-every", "5"])
|
||||
self.write(d / "learned.bin")
|
||||
self.assertEqual(engine_args(cfg)[:2], ["--expert-profile", "learned.bin"])
|
||||
self.write(d / "learned.bin", ne=256) # another model's: not used
|
||||
self.assertEqual(engine_args(cfg)[:2], ["--expert-profile", "base.bin"])
|
||||
(d / "learned.bin").write_bytes(b"STRP" + bytes(20)) # not a whole profile
|
||||
self.assertEqual(engine_args(cfg)[:2], ["--expert-profile", "base.bin"])
|
||||
self.write(d / "learned.bin")
|
||||
(d / "learned.bin").write_bytes((d / "learned.bin").read_bytes()[:30]) # truncated
|
||||
self.assertEqual(engine_args(cfg)[:2], ["--expert-profile", "base.bin"])
|
||||
|
||||
def test_no_profile_in_the_args(self):
|
||||
cfg = {"args": ["--native", "x"], "expert_profile_save": "learned.bin"}
|
||||
self.assertEqual(engine_args(cfg), ["--native", "x", "--expert-profile-save", "learned.bin"])
|
||||
|
||||
|
||||
class DraftHeadHint(unittest.TestCase):
|
||||
"""#474: a start that stopped at "the draft head does not fit" says what to change, from this start's log lines."""
|
||||
|
||||
def log(self, text, before=""):
|
||||
d = tempfile.mkdtemp()
|
||||
p = Path(d) / "engine.log"
|
||||
p.write_text(before + text, encoding="utf-8")
|
||||
return str(p), len(before.encode())
|
||||
|
||||
def test_the_engines_hint_is_relayed(self):
|
||||
p, off = self.log("strata mtp: the draft head over 106299 tokens needs 348 MiB of VRAM and 120 MiB is free.\n"
|
||||
"strata mtp: hint: a smaller draft vocabulary needs less VRAM: --draft-vocab en (...)\n"
|
||||
"strata serve: mtp: the draft head does not fit\n")
|
||||
h = start_failure_hint(p, off)
|
||||
self.assertIn("the draft head does not fit", h)
|
||||
self.assertIn("348 MiB", h)
|
||||
self.assertIn("--draft-vocab en", h)
|
||||
|
||||
def test_an_older_engine_gets_the_advice_in_words(self):
|
||||
p, off = self.log("strata serve: mtp: the draft head does not fit\n")
|
||||
self.assertIn("--draft-vocab en", start_failure_hint(p, off))
|
||||
|
||||
def test_other_failures_and_earlier_starts_add_nothing(self):
|
||||
p, off = self.log("strata serve: cannot open the pack\n")
|
||||
self.assertEqual(start_failure_hint(p, off), "")
|
||||
# an earlier start's failure (before this start's offset) is not this one's
|
||||
p, off = self.log("strata serve: cannot open the pack\n",
|
||||
before="strata serve: mtp: the draft head does not fit\n")
|
||||
self.assertEqual(start_failure_hint(p, off), "")
|
||||
self.assertEqual(start_failure_hint(None, 0), "")
|
||||
self.assertEqual(start_failure_hint(str(Path(tempfile.mkdtemp()) / "missing.log"), 0), "")
|
||||
|
||||
|
||||
class StartFailureLog(unittest.TestCase):
|
||||
"""#496: whatever stopped the engine before READY, the error carries this start's last log lines."""
|
||||
|
||||
def test_the_last_lines_of_this_start(self):
|
||||
from serve.server import start_log_tail
|
||||
d = tempfile.mkdtemp()
|
||||
p = Path(d) / "engine.log"
|
||||
before = "strata serve: an earlier start's line\n"
|
||||
p.write_text(before + "".join(f"strata serve: line {i}\n" for i in range(30)) + "\n"
|
||||
"strata serve: cannot open the pack\n", encoding="utf-8")
|
||||
tail = start_log_tail(str(p), len(before.encode()))
|
||||
self.assertIn("the engine log's last lines:", tail)
|
||||
self.assertIn("cannot open the pack", tail)
|
||||
self.assertIn("line 29", tail)
|
||||
self.assertNotIn("line 10\n", tail + "\n") # 20 lines: 11..29 and the last one
|
||||
self.assertIn("line 11", tail)
|
||||
self.assertNotIn("earlier start", tail)
|
||||
short = start_log_tail(str(p), len(before.encode()), n=3)
|
||||
self.assertEqual(short.count("\n "), 3)
|
||||
self.assertEqual(start_log_tail(str(p), p.stat().st_size), "") # nothing from this start
|
||||
self.assertEqual(start_log_tail(None, 0), "")
|
||||
self.assertEqual(start_log_tail(str(Path(d) / "missing.log"), 0), "")
|
||||
|
||||
def test_the_start_error_has_them(self):
|
||||
import serve.server as server
|
||||
fake = "import sys\nsys.stderr.write('strata serve: cannot open the pack packs/x\\n')\nsys.exit(2)\n"
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
script, log = Path(d) / "fake_strata.py", Path(d) / "strata.log"
|
||||
script.write_text(fake, encoding="utf-8")
|
||||
log.write_text("an earlier start\n", encoding="utf-8")
|
||||
real = server.subprocess.Popen
|
||||
with mock.patch.object(server.subprocess, "Popen",
|
||||
lambda cmd, **kw: real([sys.executable, str(script), *cmd[1:]], **kw)), \
|
||||
mock.patch.object(server, "narrate_start", lambda *a, **k: None):
|
||||
with self.assertRaises(RuntimeError) as cm:
|
||||
StrataEngine("strata", [], log=str(log))
|
||||
text = str(cm.exception)
|
||||
self.assertIn("exited before it was ready", text)
|
||||
self.assertIn("cannot open the pack packs/x", text)
|
||||
self.assertNotIn("an earlier start", text)
|
||||
|
||||
|
||||
class StartNarrator(unittest.TestCase):
|
||||
"""#505: the start's words say what happens to the experts - --mmap-experts loads nothing into RAM up front."""
|
||||
|
||||
def test_the_words_follow_the_flags(self):
|
||||
from serve.server import experts_loading_words
|
||||
self.assertIn("loading the experts into RAM (about 38 GB)", experts_loading_words([], "about 38 GB"))
|
||||
mapped = experts_loading_words(["--mmap-experts"], "about 47 GB")
|
||||
self.assertIn("mapping the experts from the model files (about 47 GB", mapped)
|
||||
self.assertIn("not loaded into RAM", mapped)
|
||||
self.assertIn("the GPU does not hold", experts_loading_words(["--resident-experts"], "about 47 GB"))
|
||||
budget = experts_loading_words(["--mmap-experts", "--resident-budget-gib", "71"], "about 50 GB")
|
||||
self.assertIn("up to 71 GiB", budget)
|
||||
self.assertNotIn("about 50 GB", budget)
|
||||
|
||||
def narrate(self, args, lines):
|
||||
import contextlib
|
||||
import io
|
||||
from serve.server import narrate_start
|
||||
d = tempfile.mkdtemp()
|
||||
log = Path(d) / "engine.log"
|
||||
log.write_text("".join(x + "\n" for x in lines), encoding="utf-8")
|
||||
done, out = threading.Event(), io.StringIO()
|
||||
with contextlib.redirect_stdout(out):
|
||||
t = threading.Thread(target=narrate_start, args=(str(log), 0, args, done))
|
||||
t.start()
|
||||
time.sleep(0.8)
|
||||
done.set()
|
||||
t.join(5)
|
||||
return out.getvalue()
|
||||
|
||||
def test_a_mapped_start(self):
|
||||
said = self.narrate(["--mmap-experts"], ["strata generate: experts via mmap (--mmap-experts; the GGUF shards "
|
||||
"in place, no experts.bin)"])
|
||||
self.assertIn("mapping the experts from the model files", said)
|
||||
self.assertNotIn("loading the experts into RAM", said)
|
||||
|
||||
def test_an_arena_start(self):
|
||||
said = self.narrate([], ["strata generate: expert arena: resident, 31.64 GiB",
|
||||
"strata generate: loaded 31.64 GiB at 3.17 GiB/s"])
|
||||
self.assertIn("loading the experts into RAM (tens of GB)", said)
|
||||
self.assertIn("experts loaded: 31.64 GiB at 3.17 GiB/s", said)
|
||||
|
||||
|
||||
class CancelledRead(unittest.TestCase):
|
||||
"""#471: a request cancelled while its prompt was read is recorded with the tokens the engine read (the DONE
|
||||
line's 15th field), not the whole prompt; an older engine's line (no such field) keeps the whole prompt."""
|
||||
|
||||
def test_parse_done_read_field(self):
|
||||
e = SimpleNamespace()
|
||||
StrataEngine._parse_done(e, "DONE 0 98179 17565.0 0.0 cancel 0 0 0 0 0 0 0 0.0 12288")
|
||||
self.assertEqual((e.last["prompt_tokens"], e.last["prompt_read"], e.last["finish"]), (98179, 12288, "cancel"))
|
||||
StrataEngine._parse_done(e, "DONE 0 98179 17565.0 0.0 cancel 0 0 0 0 0 0 0 0.0")
|
||||
self.assertNotIn("prompt_read", e.last)
|
||||
|
||||
def test_prompt_tokens_seen(self):
|
||||
last = {"finish": "cancel", "reused": 1000, "prompt_read": 2000}
|
||||
self.assertEqual(prompt_tokens_seen(98179, last), 3000)
|
||||
self.assertEqual(prompt_tokens_seen(98179, {**last, "finish": "stop"}), 98179) # read in full: all of it
|
||||
self.assertEqual(prompt_tokens_seen(98179, {"finish": "cancel", "reused": 0}), 98179) # an older engine
|
||||
self.assertEqual(prompt_tokens_seen(98179, {}), 98179) # no DONE at all
|
||||
self.assertEqual(prompt_tokens_seen(10, {**last, "prompt_read": 50}), 10) # never past the prompt
|
||||
|
||||
def test_history_and_totals(self):
|
||||
tok = ByteTokenizer()
|
||||
engine = DoneLineEngine(tok, "</think>\n\nok", max_context=CTX, done_lines=[
|
||||
"DONE 0 20 400.0 0.0 cancel 0 0 0 0 0 0 0 0.0 8", # stopped after 8 prompt tokens
|
||||
"DONE 4 20 40.0 30.0 stop 0 0 0 0 0 0 0 0.0 20",
|
||||
"DONE 0 20 400.0 0.0 cancel 0 0 0"]) # an older engine: no read count
|
||||
svc = Service(engine, tok, ChatTemplate(ROOT / "serve/chat_template.jinja"))
|
||||
httpd = serve(svc, port=0)
|
||||
base = f"http://127.0.0.1:{httpd.server_address[1]}"
|
||||
try:
|
||||
for _ in range(3):
|
||||
body = json.dumps({"model": "m", "max_tokens": 10, "messages": [{"role": "user", "content": "hi"}]})
|
||||
with urllib.request.urlopen(urllib.request.Request(base + "/v1/chat/completions", data=body.encode(),
|
||||
headers={"Content-Type": "application/json"}),
|
||||
timeout=30) as r:
|
||||
self.assertEqual(r.status, 200)
|
||||
with urllib.request.urlopen(base + "/metrics", timeout=10) as r:
|
||||
m = json.loads(r.read())
|
||||
finally:
|
||||
httpd.shutdown()
|
||||
httpd.server_close()
|
||||
rows = m["requests"][::-1] # oldest first
|
||||
total = rows[0]["prompt_total"]
|
||||
self.assertGreater(total, 8)
|
||||
self.assertEqual([r["prompt_total"] for r in rows], [total] * 3)
|
||||
self.assertEqual([(r["prompt_tokens"], r["prompt_read"]) for r in rows], [(8, 8), (total, 20), (total, None)])
|
||||
self.assertEqual(m["totals"]["prompt_tokens"], 8 + 2 * total)
|
||||
|
||||
|
||||
class LiveRate(unittest.TestCase):
|
||||
"""The Monitor's Speed readout: live.tok_s is a rate, and a request that never got a DONE keeps no counters.
|
||||
|
||||
@@ -1828,5 +2128,163 @@ class AmdTelemetry(unittest.TestCase):
|
||||
self.assertEqual(svc.free_vram_mib(), 26 << 10)
|
||||
|
||||
|
||||
class SilentEngine(unittest.TestCase):
|
||||
"""#481: an engine that prints nothing for engine_silence_s during a request (or never acknowledges a STOP) has
|
||||
lost step with the server: it is ended and the request fails with EngineDied, instead of waiting forever."""
|
||||
|
||||
def bare(self, silence, can_stop=False):
|
||||
import io
|
||||
import queue
|
||||
engine = StrataEngine.__new__(StrataEngine)
|
||||
engine.proc = mock.Mock()
|
||||
engine.proc.stdin = io.StringIO()
|
||||
engine.proc.poll.return_value = None
|
||||
engine.lines, engine.can_stop, engine.max_context = queue.Queue(), can_stop, 4096
|
||||
engine.silence_s, engine.log_path = silence, None
|
||||
return engine
|
||||
|
||||
def later(self, engine, delay, *lines):
|
||||
def put():
|
||||
time.sleep(delay)
|
||||
for x in lines:
|
||||
engine.lines.put(x)
|
||||
threading.Thread(target=put, daemon=True).start()
|
||||
|
||||
def test_silence_mid_answer_ends_the_engine(self):
|
||||
from serve.server import EngineSilent
|
||||
engine = self.bare(0.3)
|
||||
engine.lines.put("T 5")
|
||||
gen = engine.generate([1], 10, {}, threading.Event())
|
||||
self.assertEqual(next(gen), 5)
|
||||
t0 = time.monotonic()
|
||||
with self.assertRaises(EngineSilent) as cm:
|
||||
next(gen)
|
||||
self.assertIsInstance(cm.exception, EngineDied) # every EngineDied path handles it
|
||||
self.assertLess(time.monotonic() - t0, 5)
|
||||
engine.proc.kill.assert_called_once()
|
||||
self.assertFalse(engine.alive()) # the next request restarts it
|
||||
self.assertIn("#481", engine.death_note())
|
||||
self.assertNotIn("STOP", engine.proc.stdin.getvalue()) # nothing is listening: no STOP, no drain
|
||||
|
||||
def test_prompt_chunks_set_the_wait(self):
|
||||
# a PP line every second, at 100 tok/s: far over a 0.3 s silence, but each chunk is on time for its size
|
||||
engine = self.bare(0.3)
|
||||
engine.lines.put("RESUME 0")
|
||||
engine.lines.put("PP 100 300 1000 100.0")
|
||||
self.later(engine, 1.0, "PP 200 300 2000 100.0", "T 7", "DONE 1 300 2000 1 length")
|
||||
self.assertEqual(list(engine.generate([1], 10, {}, threading.Event())), [None, None, 7])
|
||||
engine.proc.kill.assert_not_called()
|
||||
|
||||
def test_a_long_first_chunk_is_allowed(self):
|
||||
# 100 prompt tokens at the slowest prompt reading (50 tok/s): 2 s on top of the silence before the first PP
|
||||
engine = self.bare(0.3)
|
||||
self.later(engine, 1.0, "PP 100 100 1000 100.0", "DONE 0 100 1000 0 length")
|
||||
self.assertEqual(list(engine.generate([1] * 100, 10, {}, threading.Event())), [None])
|
||||
engine.proc.kill.assert_not_called()
|
||||
|
||||
def test_a_stop_never_acknowledged(self):
|
||||
from serve.server import EngineSilent
|
||||
engine = self.bare(0.3, can_stop=True)
|
||||
engine.lines.put("T 5")
|
||||
gen = engine.generate([1], 10, {}, threading.Event())
|
||||
self.assertEqual(next(gen), 5)
|
||||
with self.assertRaises(EngineSilent):
|
||||
gen.close() # the consumer stopped: STOP, then the drain
|
||||
self.assertIn("STOP", engine.proc.stdin.getvalue())
|
||||
engine.proc.kill.assert_called_once()
|
||||
self.assertFalse(engine.alive())
|
||||
|
||||
def test_zero_waits_as_before(self):
|
||||
engine = self.bare(0)
|
||||
self.later(engine, 0.5, "T 5", "DONE 1 1 1 1 length")
|
||||
self.assertEqual(list(engine.generate([1], 10, {}, threading.Event())), [5])
|
||||
|
||||
def test_config(self):
|
||||
from serve.server import ENGINE_SILENCE_S, engine_silence_s
|
||||
self.assertEqual(engine_silence_s({}), ENGINE_SILENCE_S)
|
||||
self.assertEqual(engine_silence_s({"engine_silence_s": 0}), 0.0)
|
||||
self.assertEqual(engine_silence_s({"engine_silence_s": 900}), 900.0)
|
||||
for bad in (-1, "300", True):
|
||||
with self.assertRaises(ValueError):
|
||||
engine_silence_s({"engine_silence_s": bad})
|
||||
|
||||
|
||||
FAKE_LOST_STEP = '''import pathlib, sys, time
|
||||
mark = pathlib.Path(sys.argv[sys.argv.index("--mark") + 1])
|
||||
mode = sys.argv[sys.argv.index("--mode") + 1]
|
||||
print("READY 4096 stop", flush=True)
|
||||
for line in sys.stdin:
|
||||
if line.startswith("QUIT"):
|
||||
break
|
||||
if line.startswith("GEN"):
|
||||
if not mark.exists(): # the first engine loses step (#481): it never says DONE
|
||||
mark.touch()
|
||||
print("T 104", flush=True)
|
||||
if mode == "stop":
|
||||
print("T 257", flush=True) # <|im_end|>: the server STOPs and drains, and the engine is silent
|
||||
time.sleep(3600)
|
||||
print("T 111", flush=True)
|
||||
print("T 107", flush=True)
|
||||
print("DONE 2 5 1.0 1.0 length", flush=True)
|
||||
'''
|
||||
|
||||
|
||||
class LostStep(unittest.TestCase):
|
||||
"""#481 over HTTP, with a real process: the request with the silent engine ends with an error, and the next one
|
||||
starts the engine again and is answered (the server used to wait forever, holding the request FIFO)."""
|
||||
|
||||
def run_mode(self, mode, stream):
|
||||
import serve.server as server
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
script, mark = Path(d) / "fake_strata.py", Path(d) / "lost"
|
||||
script.write_text(FAKE_LOST_STEP, encoding="utf-8")
|
||||
real = server.subprocess.Popen
|
||||
with mock.patch.object(server.subprocess, "Popen",
|
||||
lambda cmd, **kw: real([sys.executable, str(script), *cmd[1:]], **kw)):
|
||||
eng = StrataEngine("strata", ["--mark", str(mark), "--mode", mode])
|
||||
eng.silence_s = 0.5
|
||||
tok = ByteTokenizer()
|
||||
svc = Service(eng, tok, ChatTemplate(ROOT / "serve/chat_template.jinja"))
|
||||
httpd = serve(svc, port=0)
|
||||
base = f"http://127.0.0.1:{httpd.server_address[1]}"
|
||||
body = {"model": "m", "max_tokens": 2, "reasoning_effort": "none", "stream": stream,
|
||||
"messages": [{"role": "user", "content": "hi"}]}
|
||||
try:
|
||||
first = eng.proc
|
||||
req = urllib.request.Request(base + "/v1/chat/completions", data=json.dumps(body).encode(),
|
||||
headers={"Content-Type": "application/json"})
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=60) as r:
|
||||
text = r.read().decode()
|
||||
except urllib.error.HTTPError as e:
|
||||
self.assertEqual(e.code, 503)
|
||||
text = e.read().decode()
|
||||
self.assertIn("the next request restarts it", text)
|
||||
self.assertIsNotNone(first.poll(), "the silent engine still runs")
|
||||
self.assertEqual(svc.history[-1]["finish"], "error")
|
||||
with urllib.request.urlopen(req, timeout=60) as r:
|
||||
text = r.read().decode()
|
||||
self.assertIsNot(eng.proc, first)
|
||||
if stream:
|
||||
answer = "".join(json.loads(x[6:])["choices"][0]["delta"].get("content") or ""
|
||||
for x in text.splitlines() if x.startswith("data: {") and "choices" in x)
|
||||
else:
|
||||
answer = json.loads(text)["choices"][0]["message"]["content"]
|
||||
self.assertEqual(answer, "ok")
|
||||
finally:
|
||||
httpd.shutdown()
|
||||
httpd.server_close()
|
||||
eng.unload()
|
||||
|
||||
def test_silent_mid_answer(self):
|
||||
self.run_mode("silent", stream=False)
|
||||
|
||||
def test_silent_mid_stream(self):
|
||||
self.run_mode("silent", stream=True)
|
||||
|
||||
def test_stop_never_acknowledged(self):
|
||||
self.run_mode("stop", stream=False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -69,14 +69,23 @@ HF_REVISIONS = {
|
||||
}
|
||||
|
||||
|
||||
HF_DEFAULT = "https://huggingface.co"
|
||||
|
||||
|
||||
def hf_endpoint() -> str:
|
||||
"""#495: the Hugging Face host - HF_ENDPOINT as huggingface_hub reads it (a mirror, e.g. https://hf-mirror.com),
|
||||
else huggingface.co. The pinned revisions and the SHA-256 checks are the same whichever host serves the files."""
|
||||
return (os.environ.get("HF_ENDPOINT") or "").strip().rstrip("/") or HF_DEFAULT
|
||||
|
||||
|
||||
def hf(repo: str) -> str:
|
||||
"""The download folder of a Hugging Face repository at its pinned revision."""
|
||||
return f"https://huggingface.co/{repo}/resolve/{HF_REVISIONS[repo]}/"
|
||||
return f"{hf_endpoint()}/{repo}/resolve/{HF_REVISIONS[repo]}/"
|
||||
|
||||
|
||||
def hf_unpinned(url: str) -> str:
|
||||
"""The same file at the repository's current revision (main)."""
|
||||
return re.sub(r"(https://huggingface\.co/.+?/resolve/)[0-9a-f]{40}/", r"\1main/", url, count=1)
|
||||
return re.sub(r"^(https?://[^/]+/.+?/resolve/)[0-9a-f]{40}/", r"\1main/", url, count=1)
|
||||
|
||||
|
||||
HF = hf("ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF")
|
||||
@@ -94,7 +103,7 @@ 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
|
||||
MIN_ENGINE = (0, 1, 34) # v0.1.34: AMD on Windows (a ready-made HIP engine), an MCP server for AI assistants (tools/strata_mcp.py), a shorter README; v0.1.33: a portable image encoder again (#411 #412), setup recommends instead of forcing (#406 #403 #364 #384), fixes #352 #365 #369 #371 #375 #393 #408 #414; v0.1.32: split prompts faster (#340), AMD router +12%, Unsloth Q4 in setup, faster Q4 prompts, #326/#327/#342/#344 fixes, PR batch; v0.1.31: Unsloth UD-Q4_K_XL (experimental), GGUF-in-place low-RAM mode, Windows GGUF load 2x, server race + tokenizer fixes, AMD intrinsics; v0.1.30: short prompts faster (streaming from 1024 tokens), resident low-RAM variant, multi-GPU session carve, RDNA4; v0.1.29: sampled answers faster (split top-k), #154 correctness fixes; v0.1.28: the expert cache reserves the draft head, a cancelled request no longer fails the next; v0.1.27: RTX 20 (sm_75) in the ready-made engine, the HIP build without CUDA headers; v0.1.26: the draft layer's prompt pass in batches; v0.1.25: faster prompts (grouping off the copy engine, fused hyper-connection kernels), AMD HIP backend, --kv k8v4; v0.1.24: long prompts faster (QSA select on tensor cores); v0.1.23: image requests honor sampling, 8 GB cards start, batched verify window; v0.1.22: faster prompts (tensor-core attention), multi-GPU across images/steering/KV streaming; v0.1.21: multi-GPU layer split (--gpus); v0.1.20: system-prompt checkpoint, PCIe probe, hit rate; v0.1.19: penalties
|
||||
MIN_ENGINE = (0, 1, 37) # v0.1.37: a silent engine is restarted (#481), Windows AMD counts the desktop's VRAM (#380 #377 #497), a steadier PCIe probe (#485), fixes #496 #495 #498 #505 #493; v0.1.36: a cancelled prompt logged as read so far (#471), the draft-head hint (#474), UPDATE.bat (#475), --expert-profile-save (#477); v0.1.35: Windows AMD uses its bundled HIP runtime (#468 #461), the low-RAM resident mode on Windows 32 GB (#467), fixes #460 #459 #446 #447 #457 #448 #444; v0.1.34: AMD on Windows (a ready-made HIP engine), an MCP server for AI assistants (tools/strata_mcp.py), a shorter README; v0.1.33: a portable image encoder again (#411 #412), setup recommends instead of forcing (#406 #403 #364 #384), fixes #352 #365 #369 #371 #375 #393 #408 #414; v0.1.32: split prompts faster (#340), AMD router +12%, Unsloth Q4 in setup, faster Q4 prompts, #326/#327/#342/#344 fixes, PR batch; v0.1.31: Unsloth UD-Q4_K_XL (experimental), GGUF-in-place low-RAM mode, Windows GGUF load 2x, server race + tokenizer fixes, AMD intrinsics; v0.1.30: short prompts faster (streaming from 1024 tokens), resident low-RAM variant, multi-GPU session carve, RDNA4; v0.1.29: sampled answers faster (split top-k), #154 correctness fixes; v0.1.28: the expert cache reserves the draft head, a cancelled request no longer fails the next; v0.1.27: RTX 20 (sm_75) in the ready-made engine, the HIP build without CUDA headers; v0.1.26: the draft layer's prompt pass in batches; v0.1.25: faster prompts (grouping off the copy engine, fused hyper-connection kernels), AMD HIP backend, --kv k8v4; v0.1.24: long prompts faster (QSA select on tensor cores); v0.1.23: image requests honor sampling, 8 GB cards start, batched verify window; v0.1.22: faster prompts (tensor-core attention), multi-GPU across images/steering/KV streaming; v0.1.21: multi-GPU layer split (--gpus); v0.1.20: system-prompt checkpoint, PCIe probe, hit rate; v0.1.19: penalties
|
||||
PY_PACKAGES = ["numpy", "jinja2", "regex", "pyyaml", "tqdm", "requests", "cmake", "ninja", "pillow", "psutil"]
|
||||
REQUIREMENTS = ROOT / "requirements.txt" # the same packages and their dependencies, pinned (#214)
|
||||
|
||||
@@ -390,6 +399,10 @@ def gpus():
|
||||
GPU_PICK = None # --gpu N (issue #51); None: the card with the most VRAM
|
||||
SPLIT_MIN_VRAM_GB = 8 # a card sharing a model holds the dense weights and its own
|
||||
# prompt buffers too (docs/MULTI_GPU.md)
|
||||
SPLIT_PROMPT_VRAM_GB = 12 # #448: a split stage lends a prompt chunk's buffers from its
|
||||
# own cache; below this one cannot fund a 4096-token chunk
|
||||
# (a 10 GB RTX 3080 beside a 32 GB card: 512 tokens, prompts
|
||||
# 6.2x slower), while the big card alone could
|
||||
|
||||
|
||||
def cc(g) -> str:
|
||||
@@ -441,6 +454,20 @@ def together_ok(found) -> list:
|
||||
return ok_ if len(ok_) >= 2 else []
|
||||
|
||||
|
||||
def split_short(cards) -> list:
|
||||
"""#448: the later cards of a split that would cap its prompt chunk below what the first card alone reads (a card
|
||||
under SPLIT_PROMPT_VRAM_GB beside one that has it). Empty: the split is recommended as before."""
|
||||
if len(cards) < 2 or cards[0]["vram_gb"] < SPLIT_PROMPT_VRAM_GB - 0.5:
|
||||
return []
|
||||
return [g for g in cards[1:] if g["vram_gb"] < SPLIT_PROMPT_VRAM_GB - 0.5]
|
||||
|
||||
|
||||
def split_short_note(g) -> str:
|
||||
return (f"GPU {g['index']} ({g['name']}, {g['vram_gb']:.0f} GB) is too small to lend a split its prompt buffers: "
|
||||
"it would cap prompt reading at 512-2048-token chunks, several times slower than the first card alone "
|
||||
"(#448). It can serve as a helper expert cache instead (docs/SECOND_GPU.md)")
|
||||
|
||||
|
||||
def parse_gpus(text, found) -> list:
|
||||
"""--gpus / --gpu with several: "0,2" or "all" (every card that can share the model)."""
|
||||
if str(text).strip().lower() == "all":
|
||||
@@ -549,13 +576,18 @@ def choose_gpus(a, found) -> list:
|
||||
say(" experts of its own layers, so together they hold about twice as many, and prompts are read about 20%")
|
||||
say(" faster (details: docs/MULTI_GPU.md). A much slower extra card can also make it slower.")
|
||||
opts = [can[:2]] + ([can] if len(can) > 2 else []) + [[g] for g in single]
|
||||
# #448: a pair whose second card cannot lend a 4096-token chunk recommends the first card alone (still offered)
|
||||
short = split_short(can[:2])
|
||||
rec = next(i for i, o in enumerate(opts, 1) if o == [can[0]]) if short else 1
|
||||
for i, o in enumerate(opts, 1):
|
||||
label = (" + ".join(gpu_name(g) for g in o) + " together") if len(o) > 1 else gpu_name(o[0]) + " only"
|
||||
say(f" {i}) {label}" + (" (recommended)" if i == 1 else ""))
|
||||
say(f" {i}) {label}" + (" (recommended)" if i == rec else ""))
|
||||
for g in found:
|
||||
if gpu_problem(g, together=True) is not None:
|
||||
say(f" (GPU {g['index']}, {g['name']}: {gpu_problem(g, together=True)})")
|
||||
pick = opts[int(ask("Which GPUs?", [str(i) for i in range(1, len(opts) + 1)], "1", a.yes or a.check)) - 1]
|
||||
for g in short:
|
||||
say(f" ({split_short_note(g)})")
|
||||
pick = opts[int(ask("Which GPUs?", [str(i) for i in range(1, len(opts) + 1)], str(rec), a.yes or a.check)) - 1]
|
||||
return [g["index"] for g in pick]
|
||||
|
||||
|
||||
@@ -575,6 +607,33 @@ def split_mmap(cfg: dict) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
def unsloth_split_need_gb(model="UD-Q4_K_XL") -> float:
|
||||
"""#498: the RAM UD-Q4_K_XL needs on several GPUs, where it has no RAM budget (the engine refuses
|
||||
--resident-budget-gib with a layer split): its GGUF files and UNSLOTH_RAM_LEFT_GB more (~135 GB). Measured safe
|
||||
at 165 GiB (2x RTX 3090: MemAvailable never under 68 GiB); the 0-free case of #384 was 47 GB with a 70 GB model."""
|
||||
return MODELS[model]["download_gb"] + UNSLOTH_RAM_LEFT_GB
|
||||
|
||||
|
||||
def split_budget(cfg: dict) -> bool:
|
||||
"""#498: a UD-Q4_K_XL config (its RAM budget, --resident-budget-gib) started on several GPUs. The engine refuses
|
||||
the budget with a layer split (it exited with code 2), so the split runs without it - all the experts loaded into
|
||||
RAM at start - where the RAM holds the GGUFs and 24 GB more; else setup stops, before the config is saved. True
|
||||
when the config changed."""
|
||||
a = cfg.get("args", [])
|
||||
if "--resident-budget-gib" not in a:
|
||||
return False
|
||||
need, ram = unsloth_split_need_gb(), ram_gb()
|
||||
if ram < need:
|
||||
fail(f"UD-Q4_K_XL cannot share its RAM budget across GPUs (the engine has no layer split with it), and without "
|
||||
f"the budget it needs ~{need:.0f} GB of RAM (its GGUF files and {UNSLOTH_RAM_LEFT_GB} GB more); this PC "
|
||||
f"has {ram:.0f} GB", "start it on one GPU: START-HERE.bat --gpu N (Linux: ./setup.sh --gpu N)")
|
||||
i = a.index("--resident-budget-gib")
|
||||
del a[i:i + 2]
|
||||
ok("UD-Q4_K_XL on several GPUs: no RAM budget (the engine has none with a layer split) - all its experts are "
|
||||
"loaded into RAM from the model files at start, and the files pass through the OS file cache (#498)")
|
||||
return True
|
||||
|
||||
|
||||
def offer_together(cfg_path: Path, cfg: dict, yes: bool) -> dict:
|
||||
"""Starting a model set up for one card on a PC with two or more that can share it: asked once (the answer is
|
||||
saved in its config)."""
|
||||
@@ -584,6 +643,10 @@ def offer_together(cfg_path: Path, cfg: dict, yes: bool) -> dict:
|
||||
can = together_ok(found)
|
||||
if not can:
|
||||
return cfg
|
||||
# #498: UD-Q4_K_XL's RAM budget has no layer split; without it the RAM must hold the GGUFs and 24 GB more
|
||||
budget = "--resident-budget-gib" in cfg.get("args", [])
|
||||
if budget and ram_gb() < unsloth_split_need_gb():
|
||||
return cfg
|
||||
pair = can[:2]
|
||||
cfg["gpus_asked"] = True
|
||||
# #364 #384: the resident low-RAM variant stays on one card unless the user says otherwise (its RAM use is steady)
|
||||
@@ -595,19 +658,27 @@ def offer_together(cfg_path: Path, cfg: dict, yes: bool) -> dict:
|
||||
say(" This model runs in the low-RAM mode with its experts kept in RAM, on one GPU (recommended: steady RAM")
|
||||
say(" use). On both, the experts the GPUs do not hold are read through the OS file cache instead: faster in")
|
||||
say(" two reports (#364, #384), but RAM can fill up to 0 free during long prompts.")
|
||||
if budget:
|
||||
say(" This model (UD-Q4_K_XL) runs on one GPU with a RAM budget of its experts (recommended: the tested")
|
||||
say(" setup). On both it has no budget: all its experts are loaded into RAM at start, which this PC's RAM")
|
||||
say(" holds - about twice as fast in #498 (2x RTX 3090: 31 -> 64-78 tokens/s).")
|
||||
short = split_short(pair) # #448: one card recommended (asked "n" by default), as for --resident
|
||||
for g in short:
|
||||
say(f" {split_short_note(g)}.")
|
||||
missing = [g for g in pair if not engine_runs_on(g)]
|
||||
if missing:
|
||||
say(" The installed engine has no code for " + ", ".join(g["name"] for g in missing) + ": to use them "
|
||||
"together, run START-HERE.bat --setup --gpus " + ",".join(str(g["index"]) for g in pair))
|
||||
elif ask(" Use both from now on? (you can change it later: START-HERE.bat --gpu N for one card)",
|
||||
["y", "n"], "n" if resident else "y", yes) == "y":
|
||||
["y", "n"], "n" if resident or short or budget else "y", yes) == "y":
|
||||
cfg["gpu"] = [g["index"] for g in pair]
|
||||
cfg["layer_split"] = cfg.get("layer_split") or "auto"
|
||||
split_mmap(cfg)
|
||||
split_budget(cfg)
|
||||
ok("from now on this model runs on " + " + ".join(gpu_name(g) for g in pair))
|
||||
else:
|
||||
ok("staying on one GPU (START-HERE.bat --gpus " + ",".join(str(g["index"]) for g in pair) + " switches)")
|
||||
cfg_path.write_text(json.dumps(cfg, indent=1), encoding="utf-8")
|
||||
write_config(cfg_path, cfg)
|
||||
return cfg
|
||||
|
||||
|
||||
@@ -797,6 +868,55 @@ def gguf_dir_shards(folder: Path, fam: dict, model: str) -> list[Path]:
|
||||
return [first.with_name("%s-%05d-of-%05d.gguf" % (stem, i, total)) for i in range(1, total + 1)]
|
||||
|
||||
|
||||
# #444: the quantization in a GGUF's name (Unsloth's UD-IQ3_XXS, a K-quant's Q2_K_XL, a GSQ-RCO IQ3_S, ...)
|
||||
GGUF_QUANT = re.compile(r"(?<![A-Za-z0-9])((?:UD-)?(?:I?Q\d+(?:_[A-Za-z0-9]+)*|BF16|F16|F32))"
|
||||
r"(?=-\d{5}-of-\d{5}\.gguf$|\.gguf$)", re.I)
|
||||
SUPPORTED_GGUFS = ("Strata runs ISTA-DASLab's GSQ-RCO files (Qwen3.8-Flash-Next Q2_0, IQ2_XS, IQ3_XXS, IQ3_S; Swift "
|
||||
"1.5's; the Coder's IQ1_M) and Unsloth's UD-Q4_K_XL only: other GGUFs (Unsloth's UD-IQ3_XXS or "
|
||||
"UD-Q2_K_XL, K-quants) cannot be used")
|
||||
|
||||
|
||||
def gguf_unsupported(name: str) -> str | None:
|
||||
"""#444: the quantization a GGUF's name says, when it is one Strata cannot run (not a setup size); else None."""
|
||||
m = GGUF_QUANT.search(name)
|
||||
return m.group(1) if m and m.group(1).upper() not in MODELS and not name.lower().startswith("mmproj") else None
|
||||
|
||||
|
||||
def gguf_choice(name: str) -> tuple | None:
|
||||
"""#444: (--family, --model) whose published first shard this file is, or None (the Coder's IQ1_M is named like
|
||||
the original's sizes: the size tells them apart)."""
|
||||
for f, d in FAMILIES.items():
|
||||
for m in MODELS:
|
||||
if f in MODELS[m].get("families", ("qwen", "swift")) and name == d["file"].format(q=m, i=1):
|
||||
return f, m
|
||||
return None
|
||||
|
||||
|
||||
def gguf_dir_problem(folder: Path, first: Path, fam: dict, model: str) -> tuple | None:
|
||||
"""#444: (message, hint) when --gguf-dir has no file setup can use for this choice: the chosen shard is a GGUF
|
||||
Strata cannot run (Unsloth's UD-IQ3_XXS taken for IQ3_XXS by its name), or it is missing and the folder holds
|
||||
other GGUFs - then the hint names the --family/--model of the usable ones. None otherwise (a missing shard in a
|
||||
folder without GGUFs stays check_shards' "missing")."""
|
||||
bad = gguf_unsupported(first.name) if first.exists() else None
|
||||
if bad:
|
||||
return f"{first.name} is {bad}, a GGUF Strata cannot run", SUPPORTED_GGUFS
|
||||
if first.exists():
|
||||
return None
|
||||
firsts = sorted(p.name for p in folder.glob("*.gguf")
|
||||
if not p.name.lower().startswith("mmproj") and (not SHARD_NAME.search(p.name)
|
||||
or SHARD_NAME.search(p.name).group(1) == "00001"))
|
||||
usable = list(dict.fromkeys(c for c in map(gguf_choice, firsts) if c))
|
||||
unusable = [n for n in firsts if gguf_unsupported(n)]
|
||||
if not usable and not unusable:
|
||||
return None
|
||||
hint = SUPPORTED_GGUFS
|
||||
if unusable:
|
||||
hint += ".\n Not usable here: " + ", ".join(unusable)
|
||||
if usable:
|
||||
hint += ".\n Usable here: " + ", ".join(f"--family {f} --model {m}" for f, m in usable)
|
||||
return f"{folder} has no {fam['title']} {model} file", hint
|
||||
|
||||
|
||||
def verify_sha256(s: Path, size: int, sha: str) -> None:
|
||||
"""A shard's size and SHA-256 against the pinned values (the Unsloth file); the result is kept in its finish mark,
|
||||
so the ~5 minutes of hashing 111 GB happen once. A wrong file is deleted, so the next run downloads it again."""
|
||||
@@ -1192,6 +1312,7 @@ def hip_devices(probe: Path | None = None, text: str | None = None) -> list[dict
|
||||
if not probe.exists() or not hip_engine:
|
||||
return None
|
||||
try:
|
||||
hip_runtime_beside_exe(probe.parent) # #468 #461: not the driver's System32 copy
|
||||
env = dict(os.environ) # the ready-made engine's ROCm DLLs (rocm/bin beside it)
|
||||
env["PATH"] = os.pathsep.join([str(d) for d in hip_lib_dirs(probe.parent)] + [env.get("PATH", "")])
|
||||
r = subprocess.run([str(probe), "--list-devices"], capture_output=True, text=True, timeout=120,
|
||||
@@ -1232,6 +1353,31 @@ def hip_lib_dirs(eng: Path) -> list[Path]:
|
||||
return [eng / d for d in rel if (eng / d).is_dir()]
|
||||
|
||||
|
||||
# #468 #461: the HIP runtime the ready-made engine was built with, next to strata.exe. Windows looks for an imported
|
||||
# DLL in the exe's folder, then System32, and only then on PATH (where rocm/bin is): an AMD driver that installs its own
|
||||
# amdhip64_7.dll in System32 won, and the bundled rocBLAS/hipBLAS ran on that runtime - an access violation (W7900) or
|
||||
# hipErrorInvalidDeviceFunction (7900 XTX) on the first prompt. Only the runtime and the compiler it loads by name:
|
||||
# rocBLAS/hipBLASLt stay in rocm/bin, where they find their kernel libraries and ../.kpack.
|
||||
HIP_RUNTIME_DLLS = ("amdhip64_*.dll", "amd_comgr*.dll")
|
||||
|
||||
|
||||
def hip_runtime_beside_exe(eng: Path) -> None:
|
||||
"""Copy the bundled HIP runtime DLLs from rocm/bin next to the engine's exes when missing or different (a 0.1.34
|
||||
install, whose zip had them in rocm/bin only, is fixed on its next start)."""
|
||||
for d in hip_lib_dirs(eng):
|
||||
for pat in HIP_RUNTIME_DLLS:
|
||||
for src in d.glob(pat):
|
||||
dst = eng / src.name
|
||||
try:
|
||||
if dst.exists() and dst.stat().st_size == src.stat().st_size and \
|
||||
dst.stat().st_mtime >= src.stat().st_mtime:
|
||||
continue
|
||||
shutil.copy2(src, dst)
|
||||
except OSError as e: # e.g. the engine is running and holds the old copy
|
||||
warn(f"could not put {src.name} next to the AMD engine ({e}); if the engine stops on its first "
|
||||
"request, close Strata and run START-HERE.bat again")
|
||||
|
||||
|
||||
def hip_match(card: dict, listed: list[dict], hip: list[dict]) -> dict | None:
|
||||
"""The HIP device that is setup's `card` (from `listed`, the display-adapter order): the k-th device of the same
|
||||
architecture, k = the card's rank among the listed cards of that architecture. None when HIP has no such card."""
|
||||
@@ -1324,6 +1470,7 @@ def get_prebuilt_hip(url_base, gpu, updating=False) -> Path | None:
|
||||
p.replace(dst)
|
||||
shutil.rmtree(tmp, ignore_errors=True)
|
||||
drop_archive(z)
|
||||
hip_runtime_beside_exe(eng) # #468 #461
|
||||
ok(f"ready-made AMD engine {meta.get('version', '')} for {', '.join(meta.get('archs', []))} "
|
||||
f"(ROCm {meta.get('rocm', '?')})")
|
||||
return eng
|
||||
@@ -1338,18 +1485,32 @@ def rocm_version(root):
|
||||
return None
|
||||
|
||||
|
||||
def rocm_dev_missing(sysroot: Path) -> list:
|
||||
"""#446: the HIP development files the engine's build needs that a system ROCm lacks (a runtime-only install has
|
||||
hipcc and libhipblas but not these, and cmake's enable_language(HIP) then fails on the hip-lang package)."""
|
||||
lang = "cmake/hip-lang/hip-lang-config.cmake" # where CMake's HIP support looks for it
|
||||
need = {"lib/" + lang: [sysroot / d / lang for d in ("lib", "lib64", "lib/x86_64-unknown-linux-gnu")],
|
||||
"include/hip/hip_runtime.h": [sysroot / "include" / "hip" / "hip_runtime.h"]}
|
||||
return [name for name, paths in need.items() if not any(p.is_file() for p in paths)]
|
||||
|
||||
|
||||
def rocm_root(archs):
|
||||
"""ROCm for compiling and running the HIP engine for `archs` (one arch or a list: the cards of a layer split):
|
||||
(root, library folders). A system ROCm 7 with hipcc and hipBLAS, else AMD's TheRock wheels (ROCM_VERSION, from the
|
||||
card family's index) installed into .venv."""
|
||||
(root, library folders). A system ROCm 7 with hipcc, hipBLAS and the HIP development files (#446), else AMD's
|
||||
TheRock wheels (ROCM_VERSION, from the card family's index) installed into .venv."""
|
||||
archs = [archs] if isinstance(archs, str) else list(archs)
|
||||
sysroot = Path(os.environ.get("ROCM_PATH") or "/opt/rocm")
|
||||
if (sysroot / "bin" / "hipcc").exists() and list((sysroot / "lib").glob("libhipblas.so*")):
|
||||
ver = rocm_version(sysroot)
|
||||
if ver is None or ver >= ROCM_SYSTEM_MIN:
|
||||
missing = rocm_dev_missing(sysroot)
|
||||
if (ver is None or ver >= ROCM_SYSTEM_MIN) and not missing:
|
||||
return sysroot, [str(sysroot / "lib")]
|
||||
warn(f"the ROCm in {sysroot} is {ver[0]}.{ver[1]}; Strata needs {ROCM_SYSTEM_MIN[0]}.{ROCM_SYSTEM_MIN[1]} or "
|
||||
"newer: using AMD's wheels in .venv instead")
|
||||
if ver is not None and ver < ROCM_SYSTEM_MIN:
|
||||
warn(f"the ROCm in {sysroot} is {ver[0]}.{ver[1]}; Strata needs {ROCM_SYSTEM_MIN[0]}.{ROCM_SYSTEM_MIN[1]} "
|
||||
"or newer: using AMD's wheels in .venv instead")
|
||||
else: # #446: a runtime-only ROCm (no -dev packages): cmake would fail
|
||||
warn(f"the ROCm in {sysroot} has no HIP development files ({', '.join(missing)}): using AMD's wheels in "
|
||||
".venv instead (or install them, e.g. AMD's amdrocm-core-dev package for your ROCm and card)")
|
||||
indexes = list(dict.fromkeys(rocm_index(a) for a in archs))
|
||||
if len(indexes) > 1: # TheRock's wheels hold one GPU family's libraries
|
||||
fail(f"cards of two GPU families ({', '.join(archs)}) need a system ROCm 7 (in /opt/rocm): AMD's Python "
|
||||
@@ -1952,6 +2113,43 @@ def low_ram_together(a, model, ram, gpu, chosen) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def unsloth_together(a, model, ram, gpu, chosen) -> bool:
|
||||
"""#498: UD-Q4_K_XL with several GPUs chosen. Its RAM budget (--resident-budget-gib) has no layer split, so a
|
||||
split runs without it: all its experts loaded into RAM from the GGUFs at start, as with the 2-3-bit models - only
|
||||
where the RAM holds the GGUF files and 24 GB more (unsloth_split_need_gb; 165 GiB, 2x RTX 3090: 31 -> 64-78
|
||||
tok/s). An explicit --gpus is honoured there; otherwise asked, one GPU by default (--yes: one GPU, as before); an
|
||||
explicit --resident-budget-gib keeps one GPU. True: all of them."""
|
||||
names = " + ".join(gpu_name(g) for g in chosen)
|
||||
need = unsloth_split_need_gb(model)
|
||||
if ram < need:
|
||||
warn(f"{model} runs on one GPU here: on several it has no RAM budget and needs ~{need:.0f} GB of RAM (its GGUF "
|
||||
f"files and {UNSLOTH_RAM_LEFT_GB} GB more), this PC has {ram:.0f} - using {gpu_name(gpu)} only")
|
||||
return False
|
||||
if a.resident_budget_gib is not None:
|
||||
warn(f"--resident-budget-gib has no layer split: {model} runs on one GPU with it - using {gpu_name(gpu)} only "
|
||||
f"(leave the budget out to use {names} together)")
|
||||
return False
|
||||
note = (f"no RAM budget - all of its experts (~{MODELS[model]['arena_gb']:.0f} GB) are loaded into RAM from the "
|
||||
f"model files at start, and the files pass through the OS file cache (needs ~{need:.0f} GB of RAM, this "
|
||||
f"PC has {ram:.0f})")
|
||||
if a.gpus:
|
||||
ok(f"{model} on {names}, as you chose (--gpus): {note}")
|
||||
return True
|
||||
if not a.yes:
|
||||
say()
|
||||
say(f" {model} can run on one GPU with a RAM budget of its experts, or on {names} together without one:")
|
||||
say(f" 1) {gpu_name(gpu)} only: the most-used experts kept in RAM, the rest read from the SSD (recommended: "
|
||||
"the tested setup)")
|
||||
say(f" 2) {names} together: {note};")
|
||||
say(" about twice as fast in #498 (2x RTX 3090: 31 -> 64-78 tokens/s)")
|
||||
if ask(f"{model}: which GPUs?", ["1", "2"], "1", a.yes) == "2":
|
||||
ok(f"{model} on {names}: {note}")
|
||||
return True
|
||||
warn(f"{model} runs on one GPU: using {gpu_name(gpu)} only (--gpus " + ",".join(str(g["index"]) for g in chosen) +
|
||||
f" shares it across {names} without the RAM budget: this PC's RAM holds it)")
|
||||
return False
|
||||
|
||||
|
||||
def confirm_risk(msg, explicit, yes, stop, hint=None, question=" Go on anyway?", default="n") -> None:
|
||||
"""The owner's rule: setup recommends, it never forces. A choice setup expects to fail or run badly is said
|
||||
plainly (msg), then asked (default n; `default` keeps an older question's own default), or with --yes taken as
|
||||
@@ -2105,7 +2303,7 @@ def repoint_config(cfg_file: Path, old: Path, new: Path) -> None:
|
||||
|
||||
new_cfg = fix(cfg)
|
||||
if new_cfg != cfg:
|
||||
cfg_file.write_text(json.dumps(new_cfg, indent=1), encoding="utf-8")
|
||||
write_config(cfg_file, new_cfg)
|
||||
|
||||
|
||||
def data_folder(requested: str | None) -> tuple:
|
||||
@@ -2162,13 +2360,39 @@ def data_folder(requested: str | None) -> tuple:
|
||||
return dest, elsewhere
|
||||
|
||||
|
||||
def write_config(path: Path, cfg: dict):
|
||||
"""A run config, written whole or not at all (#459): to a temporary file first, then moved over the old one, so
|
||||
a setup stopped half-way (a closed window, a full disk) never leaves an empty strata-*.json behind."""
|
||||
tmp = path.with_name(path.name + ".tmp")
|
||||
tmp.write_text(json.dumps(cfg, indent=1), encoding="utf-8")
|
||||
os.replace(tmp, path)
|
||||
|
||||
|
||||
def readable_config(path: Path) -> bool:
|
||||
"""#459: a config that parses as a JSON object; any other gets a one-line warning naming it."""
|
||||
text = None
|
||||
try:
|
||||
text = path.read_text(encoding="utf-8-sig")
|
||||
if isinstance(json.loads(text), dict):
|
||||
return True
|
||||
why = "not a JSON object"
|
||||
except OSError as e:
|
||||
why = e.strerror or str(e)
|
||||
except ValueError: # JSONDecodeError, or bytes that are not UTF-8
|
||||
why = "the file is empty" if text is not None and not text.strip() else "not valid JSON"
|
||||
warn(f"skipped the earlier config {path} ({why}): setting this copy up without it")
|
||||
return False
|
||||
|
||||
|
||||
def previous_config(elsewhere_first: list, settings: dict):
|
||||
"""The most recently used model config of another Strata folder on this PC, for a folder that has none yet."""
|
||||
"""The most recently used model config of another Strata folder on this PC, for a folder that has none yet. One
|
||||
that does not parse (an empty or cut-off file, #459) is skipped with a warning: the newest readable one is used,
|
||||
and with none this copy is set up as a fresh install."""
|
||||
cands = []
|
||||
for folder in [*elsewhere_first, *other_installs(settings)]:
|
||||
cands += list(folder.glob("strata-*.json"))
|
||||
cands = [c for c in dict.fromkeys(cands) if c.is_file()]
|
||||
return max(cands, key=lambda p: p.stat().st_mtime) if cands else None
|
||||
return next((c for c in sorted(cands, key=lambda p: p.stat().st_mtime, reverse=True) if readable_config(c)), None)
|
||||
|
||||
|
||||
def choices_from_config(cfg_path: Path) -> dict:
|
||||
@@ -2190,7 +2414,10 @@ def choices_from_config(cfg_path: Path) -> dict:
|
||||
"vision": ("gpu" if vis.get("gpu") else "cpu") if isinstance(vis, dict) else "none",
|
||||
"esp": ("on" if Path(esp_path).name == ESP_VECTOR.name else esp_path) if esp_path else "off",
|
||||
"host": cfg.get("host"), "api_key": cfg.get("api_key"), "port": cfg.get("port"), "gpu": cfg.get("gpu"),
|
||||
"layer_split": cfg.get("layer_split")}
|
||||
"layer_split": cfg.get("layer_split"),
|
||||
# #493: --vram-reserve-mib given at setup (images write the default 700 themselves)
|
||||
"vram_reserve_mib": int(val("--vram-reserve-mib")) if (val("--vram-reserve-mib") or "").isdigit() and (
|
||||
vis is None or int(val("--vram-reserve-mib")) != VISION["gpu"]["reserve_mib"]) else None}
|
||||
|
||||
|
||||
def find_in(roots: list, rel: str):
|
||||
@@ -2255,13 +2482,20 @@ def calibrate_config(cfg_path: Path) -> bool:
|
||||
say()
|
||||
say(" Tuning Strata for this PC: the output speed is measured with a few engine settings (the PCIe share, the")
|
||||
say(" draft depth, the CPU threads). It takes about 5-10 minutes; the PC is busy meanwhile.")
|
||||
try:
|
||||
since = os.path.getsize(cfg["log"]) if cfg.get("log") and os.path.isfile(cfg["log"]) else 0
|
||||
except OSError:
|
||||
since = 0
|
||||
try:
|
||||
res = CAL.run(cfg, say=say)
|
||||
except Exception as e: # never stops an install: the defaults stay
|
||||
warn(f"the tuning did not finish ({e}): the default settings stay")
|
||||
why = CAL.engine_error(cfg.get("log"), since) # #447: the engine's own reason, not only "see the log"
|
||||
if why:
|
||||
say(f" the engine said: {why}")
|
||||
return False
|
||||
cfg["args"] = CAL.apply(cfg["args"], res["settings"])
|
||||
cfg_path.write_text(json.dumps(cfg, indent=1), encoding="utf-8")
|
||||
write_config(cfg_path, cfg)
|
||||
st = load_settings()
|
||||
st.setdefault("calibration", {})[hardware_key(cfg)] = {"settings": res["settings"], "tok_s": res["report"].get("tok_s"),
|
||||
"date": time.strftime("%Y-%m-%d")}
|
||||
@@ -2301,21 +2535,59 @@ def upgrade_config(cfg_path: Path, cfg: dict) -> dict:
|
||||
changed = True
|
||||
ok("WSL: KV streaming off (the driver pins only about 1 GB of RAM); the KV cache stays in VRAM")
|
||||
if changed:
|
||||
cfg_path.write_text(json.dumps(cfg, indent=1), encoding="utf-8")
|
||||
write_config(cfg_path, cfg)
|
||||
return cfg
|
||||
|
||||
|
||||
def update_install(have: list, a) -> int:
|
||||
"""#475: `setup.py --update` (UPDATE.bat / update.sh, after their git pull): what a plain START-HERE.bat does to
|
||||
an install before it starts the model, without starting it - the Python packages, the ready-made engine when this
|
||||
setup needs a newer one (MIN_ENGINE; a compiled engine when its source changed), each installed model's config
|
||||
upgrades and its draft subset. No question is asked and the model files are not touched; a model still running
|
||||
keeps its engine (update_installed_engine says to close it and run this again)."""
|
||||
if not have:
|
||||
say(" No model is installed in this Strata folder yet: run START-HERE.bat (Linux: ./setup.sh) to set it up -")
|
||||
say(" it finds an earlier install's model files next to it and reuses them.")
|
||||
return 0
|
||||
pip_install(requirement_lines() if REQUIREMENTS.exists() else PY_PACKAGES,
|
||||
"numpy, jinja2, regex, pyyaml, tqdm, requests, cmake, ninja, pillow, psutil")
|
||||
if not a.build:
|
||||
update_installed_engine(a.prebuilt)
|
||||
for cfg_path in have:
|
||||
cfg = upgrade_config(cfg_path, json.loads(cfg_path.read_text(encoding="utf-8-sig")))
|
||||
if "--mtp" in cfg["args"][:-1]:
|
||||
refresh_draft_vocab(Path(cfg["args"][cfg["args"].index("--mtp") + 1]), cfg.get("draft_vocab", "cjk"))
|
||||
if cfg.get("backend") == "hip" and WIN:
|
||||
hip_runtime_beside_exe(Path(cfg["exe"]).parent) # #468 #461
|
||||
ok(f"{cfg.get('model_name', cfg_path.stem)}: up to date")
|
||||
ver = engine_version(Path(json.loads(have[0].read_text(encoding="utf-8-sig"))["exe"]))
|
||||
say()
|
||||
ok("Strata is updated" + (f" (engine {'.'.join(map(str, ver))})" if any(ver) else "") +
|
||||
". Start the model with " + ("START-HERE.bat" if WIN else "./setup.sh") + " when you want it.")
|
||||
return 0
|
||||
|
||||
|
||||
def start(cfg_path: Path, port: int | None, gpu: int | list | None = None, open_browser=True, yes=False,
|
||||
layer_split=None, keep=None) -> int:
|
||||
"""keep: settings given on this start that the model keeps from now on (--host, --api-key, --draft-vocab)."""
|
||||
"""keep: settings given on this start that the model keeps from now on (--host, --api-key, --draft-vocab,
|
||||
--vram-reserve-mib)."""
|
||||
cfg = upgrade_config(cfg_path, json.loads(cfg_path.read_text(encoding="utf-8-sig")))
|
||||
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")
|
||||
keep = {k: v for k, v in (keep or {}).items() if v is not None}
|
||||
reserve = keep.pop("vram_reserve_mib", None) # #493: an engine argument, kept in the config's args
|
||||
if reserve is not None:
|
||||
args = cfg["args"]
|
||||
if "--vram-reserve-mib" in args[:-1]:
|
||||
args[args.index("--vram-reserve-mib") + 1] = str(reserve)
|
||||
else:
|
||||
args += ["--vram-reserve-mib", str(reserve)]
|
||||
write_config(cfg_path, cfg)
|
||||
ok(f"saved for this model: {reserve} MiB of VRAM kept free for other programs (--vram-reserve-mib)")
|
||||
if keep and any(cfg.get(k) != v for k, v in keep.items()): # #179: a --host/--api-key on a start was ignored
|
||||
cfg.update(keep)
|
||||
cfg_path.write_text(json.dumps(cfg, indent=1), encoding="utf-8")
|
||||
write_config(cfg_path, cfg)
|
||||
ok("saved for this model: " + ", ".join("api key" if k == "api_key" else f"{k.replace('_', ' ')} {v}"
|
||||
for k, v in keep.items()))
|
||||
cfg_path.touch() # the most recently used model
|
||||
@@ -2324,6 +2596,8 @@ def start(cfg_path: Path, port: int | None, gpu: int | list | None = None, open_
|
||||
cmd = [sys.executable, str(ROOT / "serve" / "server.py"), "--engine", "strata", "--config", str(cfg_path),
|
||||
"--port", str(port or cfg.get("port", 8080))]
|
||||
if cfg.get("backend") == "hip": # AMD, numbered as HIP numbers them (setup's KFD order)
|
||||
if WIN:
|
||||
hip_runtime_beside_exe(Path(cfg["exe"]).parent) # #468 #461: also fixes a 0.1.34 install
|
||||
amd = amd_gpus()
|
||||
if isinstance(gpu, list): # --gpus: saved, this model runs on these cards from now on
|
||||
cards = amd_parse_gpus(",".join(str(i) for i in gpu), amd)
|
||||
@@ -2334,7 +2608,8 @@ def start(cfg_path: Path, port: int | None, gpu: int | list | None = None, open_
|
||||
"set it up for these cards: ./setup.sh --setup --backend hip --gpus " + ",".join(map(str, gpu)))
|
||||
cfg["gpu"], cfg["gpus_asked"] = gpu, True
|
||||
cfg["layer_split"] = layer_split or cfg.get("layer_split") or "auto"
|
||||
cfg_path.write_text(json.dumps(cfg, indent=1), encoding="utf-8")
|
||||
split_budget(cfg) # #498: before it is saved (it stops when the RAM is short)
|
||||
write_config(cfg_path, cfg)
|
||||
gpu = None
|
||||
elif gpu is not None:
|
||||
cmd += ["--gpu", str(gpu)]
|
||||
@@ -2357,7 +2632,8 @@ def start(cfg_path: Path, port: int | None, gpu: int | list | None = None, open_
|
||||
check_gpus(gpu, found, yes=yes, named=True)
|
||||
cfg["gpu"], cfg["gpus_asked"] = gpu, True
|
||||
cfg["layer_split"] = layer_split or cfg.get("layer_split") or "auto"
|
||||
cfg_path.write_text(json.dumps(cfg, indent=1), encoding="utf-8")
|
||||
split_budget(cfg) # #498: before it is saved (it stops when the RAM is short)
|
||||
write_config(cfg_path, cfg)
|
||||
gpu = None
|
||||
elif gpu is not None: # --gpu N: this start only, on that card
|
||||
check_gpus([gpu], found)
|
||||
@@ -2365,8 +2641,9 @@ def start(cfg_path: Path, port: int | None, gpu: int | list | None = None, open_
|
||||
else:
|
||||
cfg = offer_together(cfg_path, cfg, yes)
|
||||
use = gpu if gpu is not None else cfg.get("gpu")
|
||||
if isinstance(use, list) and split_mmap(cfg): # #364 #384: a resident low-RAM config on several GPUs
|
||||
cfg_path.write_text(json.dumps(cfg, indent=1), encoding="utf-8")
|
||||
# #364 #384: a resident low-RAM config on several GPUs; #498: a UD-Q4_K_XL config with its RAM budget (by hand)
|
||||
if isinstance(use, list) and (split_mmap(cfg) | split_budget(cfg)):
|
||||
write_config(cfg_path, cfg)
|
||||
if cfg.get("backend") == "hip":
|
||||
pass
|
||||
elif isinstance(use, list):
|
||||
@@ -2390,8 +2667,15 @@ def start(cfg_path: Path, port: int | None, gpu: int | list | None = None, open_
|
||||
pass
|
||||
say()
|
||||
say(" " + "-" * 100)
|
||||
say(f" Starting {cfg.get('model_name', 'the model')}: it loads {f'about {gb:.0f} GB' if gb >= 1 else '34-55 GB'} "
|
||||
"into RAM and locks part of it for the GPU.")
|
||||
size = f'about {gb:.0f} GB' if gb >= 1 else '34-55 GB'
|
||||
a_ = cfg["args"]
|
||||
if "--mmap-experts" in a_ and "--resident-budget-gib" not in a_ and "--resident-experts" not in a_:
|
||||
# #505: the mapped low-RAM mode loads nothing into RAM up front (the server's narrator says the same)
|
||||
say(f" Starting {cfg.get('model_name', 'the model')}: it maps {size} of experts from the model files (the OS "
|
||||
"file cache reads them).")
|
||||
else:
|
||||
say(f" Starting {cfg.get('model_name', 'the model')}: it loads {size} into RAM and locks part of it for the "
|
||||
"GPU.")
|
||||
say(" While it does, YOUR PC CAN BE SLOW OR STOP RESPONDING FOR 1-3 MINUTES (longer the first time after a")
|
||||
say(" restart). That is normal: please wait and don't close this window - the browser opens when it is ready.")
|
||||
say(" Later, closing this window stops the model.")
|
||||
@@ -2421,6 +2705,25 @@ def saved_draft_vocab(cfg_path: Path) -> str | None:
|
||||
return v if v in DRAFT_VOCABS else None
|
||||
|
||||
|
||||
DRAFT_VOCAB_MIB = {"cjk": 348, "cyrillic": 193, "en": 133} # the draft head's VRAM per subset (IQ3_S: the largest)
|
||||
SMALL_DRAFT_VRAM_GB = 14 # #474: below this the default subset's head can be what does not fit
|
||||
|
||||
|
||||
def draft_vocab_note(vram_gb: float, chosen: str | None) -> list[str]:
|
||||
"""#474: on a card under 14 GB, the default draft subset (cjk, ~348 MiB of VRAM) can be what does not fit at the
|
||||
start ("the draft head does not fit"), and the engine's expert cache gets what a smaller one leaves. Setup
|
||||
RECOMMENDS a smaller one here and changes nothing (the owner's rule, #403 #406): a subset chosen with
|
||||
--draft-vocab, or kept from an earlier install, gets no note. [] for every other case."""
|
||||
if chosen or not 0 < vram_gb < SMALL_DRAFT_VRAM_GB:
|
||||
return []
|
||||
start = "START-HERE.bat" if WIN else "./setup.sh"
|
||||
return [f"Tip for a {vram_gb:.0f} GB card: the draft layer's default token subset (with Chinese, Japanese and "
|
||||
f"Korean) needs up to ~{DRAFT_VOCAB_MIB['cjk']} MiB of VRAM.",
|
||||
f" For English and code answers, {start} --draft-vocab en needs up to ~{DRAFT_VOCAB_MIB['en']} MiB "
|
||||
f"(cyrillic: ~{DRAFT_VOCAB_MIB['cyrillic']}) and leaves the rest to the expert cache - and it is the",
|
||||
" fix when the start stops with \"the draft head does not fit\". The model keeps the choice."]
|
||||
|
||||
|
||||
def mtp_corrupt(mtp: Path, env=None) -> bool:
|
||||
"""#327: True when the MTP tensors an install fetched are not the pinned checkpoint's (tools/mtp_fetch.py verify,
|
||||
which hashes only files that changed since they last checked out). A mirror that ignored range requests left the
|
||||
@@ -2471,7 +2774,7 @@ def ensure_engine_for(cards, cfg_path: Path, cfg: dict, yes: bool) -> dict:
|
||||
build_engine({**main, "archs": archs}, vision, yes, get_llama_cpp())
|
||||
dirs = json.loads(info.read_text()).get("cuda_dirs") or []
|
||||
cfg["lib_dirs"] = dirs + [d for d in cfg.get("lib_dirs") or [] if d not in dirs]
|
||||
cfg_path.write_text(json.dumps(cfg, indent=1), encoding="utf-8")
|
||||
write_config(cfg_path, cfg)
|
||||
return cfg
|
||||
|
||||
|
||||
@@ -2567,6 +2870,9 @@ def main() -> int:
|
||||
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("--update", action="store_true",
|
||||
help="update the installed engine, Python packages and model settings as a start would, without "
|
||||
"starting the model (UPDATE.bat / update.sh run it after a git pull)")
|
||||
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)")
|
||||
@@ -2585,6 +2891,9 @@ def main() -> int:
|
||||
ap.add_argument("--resident-budget-gib", type=float, metavar="N",
|
||||
help="UD-Q4_K_XL: the GiB of its experts kept in RAM (default: the RAM less 24 GB, 40 on 64 GB; "
|
||||
"more is kept as you choose, with a note)")
|
||||
ap.add_argument("--vram-reserve-mib", type=int, metavar="N",
|
||||
help="VRAM in MiB the engine leaves free for other programs (a game, another model; the engine's "
|
||||
"default: 700); the expert cache takes that much less")
|
||||
ap.add_argument("--kv-streaming", choices=["auto", "on", "off"], default="auto",
|
||||
help="from a 64K context: keep the KV cache in RAM and only the attention's window in VRAM (more "
|
||||
"experts fit on the GPU); auto: when the RAM has room for it")
|
||||
@@ -2595,6 +2904,8 @@ def main() -> int:
|
||||
a = ap.parse_args()
|
||||
if a.resident_budget_gib is not None and not a.resident_budget_gib > 0:
|
||||
ap.error("--resident-budget-gib takes a number of GiB above 0, e.g. --resident-budget-gib 32")
|
||||
if a.vram_reserve_mib is not None and a.vram_reserve_mib < 0:
|
||||
ap.error("--vram-reserve-mib takes a number of MiB, 0 or more, e.g. --vram-reserve-mib 2048")
|
||||
if a.gpu is not None: # --gpu 0,2 means --gpus 0,2 (a user tried it: issue report)
|
||||
if "," in a.gpu:
|
||||
a.gpus, a.gpu = a.gpus or a.gpu, None
|
||||
@@ -2610,6 +2921,8 @@ def main() -> int:
|
||||
|
||||
# ---- 0. already installed: just start it
|
||||
have = installed_configs()
|
||||
if a.update: # #475: UPDATE.bat / update.sh - never starts the model
|
||||
return update_install(have, a)
|
||||
explicit = a.setup or a.model or a.family or a.check or a.no_start
|
||||
if not have and not explicit: # a new copy of Strata (an update unzipped elsewhere): set it
|
||||
prev = previous_config(elsewhere, load_settings()) # up like the last one, from the files already here
|
||||
@@ -2624,6 +2937,8 @@ def main() -> int:
|
||||
a.experimental_speed_projection = a.experimental_speed_projection or ch["esp"]
|
||||
a.host, a.api_key = a.host or ch["host"], a.api_key or ch["api_key"]
|
||||
a.port = a.port or ch["port"]
|
||||
if a.vram_reserve_mib is None: # #493: an explicit reserve set up before
|
||||
a.vram_reserve_mib = ch.get("vram_reserve_mib")
|
||||
if isinstance(ch.get("gpu"), list): # a layer split: set up across the same cards again
|
||||
a.gpus = a.gpus or ",".join(str(g) for g in ch["gpu"])
|
||||
a.layer_split = a.layer_split or ch.get("layer_split")
|
||||
@@ -2645,15 +2960,20 @@ def main() -> int:
|
||||
for i, c in enumerate(have, 1):
|
||||
say(f" {i}) {json.loads(c.read_text(encoding='utf-8-sig')).get('model_name', c.stem)}")
|
||||
pick_cfg = have[int(ask("Tune which one?", [str(i) for i in range(1, len(have) + 1)], "1", a.yes)) - 1]
|
||||
calibrate_config(pick_cfg)
|
||||
if not calibrate_config(pick_cfg): # #447: said again where it is not lost above the start
|
||||
say()
|
||||
warn("this PC is NOT tuned: the tuning failed (the reason is above); the model "
|
||||
+ ("keeps" if a.no_start else "starts with") + " the default settings")
|
||||
return 0 if a.no_start else start(pick_cfg, a.port, run_gpu, yes=a.yes, layer_split=a.layer_split,
|
||||
keep={"host": a.host, "api_key": a.api_key, "draft_vocab": a.draft_vocab})
|
||||
keep={"host": a.host, "api_key": a.api_key, "draft_vocab": a.draft_vocab,
|
||||
"vram_reserve_mib": a.vram_reserve_mib})
|
||||
if have and not (a.setup or a.model or a.family or a.check or a.no_start):
|
||||
if not a.build:
|
||||
update_installed_engine(a.prebuilt)
|
||||
if len(have) == 1:
|
||||
return start(have[0], a.port, run_gpu, yes=a.yes, layer_split=a.layer_split,
|
||||
keep={"host": a.host, "api_key": a.api_key, "draft_vocab": a.draft_vocab})
|
||||
keep={"host": a.host, "api_key": a.api_key, "draft_vocab": a.draft_vocab,
|
||||
"vram_reserve_mib": a.vram_reserve_mib})
|
||||
say()
|
||||
for i, c in enumerate(have, 1):
|
||||
say(f" {i}) {json.loads(c.read_text(encoding='utf-8-sig')).get('model_name', c.stem)}")
|
||||
@@ -2661,7 +2981,8 @@ def main() -> int:
|
||||
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], a.port, run_gpu, yes=a.yes, layer_split=a.layer_split,
|
||||
keep={"host": a.host, "api_key": a.api_key, "draft_vocab": a.draft_vocab})
|
||||
keep={"host": a.host, "api_key": a.api_key, "draft_vocab": a.draft_vocab,
|
||||
"vram_reserve_mib": a.vram_reserve_mib})
|
||||
|
||||
# ---- 1. the PC
|
||||
step(1, "checking your PC")
|
||||
@@ -2794,7 +3115,12 @@ def main() -> int:
|
||||
say()
|
||||
names = [m for m in MODELS if family in MODELS[m].get("families", ("qwen", "swift"))]
|
||||
if a.model and a.model not in names:
|
||||
fail(f"{fam['title']} has no {a.model} model file", "choose one of: " + ", ".join(names))
|
||||
# #444: say which family has that size, and (with --gguf-dir) which files Strata can run at all
|
||||
elsewhere_fams = [f for f in FAMILIES if f in MODELS[a.model].get("families", ("qwen", "swift"))]
|
||||
fail(f"{fam['title']} has no {a.model} model file", "choose one of: " + ", ".join(names)
|
||||
+ (f" (or {a.model}: " + ", ".join(f"--family {f} --model {a.model}" for f in elsewhere_fams) + ")"
|
||||
if elsewhere_fams else "")
|
||||
+ (f".\n {SUPPORTED_GGUFS}" if a.gguf_dir else ""))
|
||||
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}"
|
||||
@@ -2808,10 +3134,11 @@ def main() -> int:
|
||||
say(f" {i}) {m:8s} {d['about']}; download {d['download_gb']:.0f} GB, uses ~{d['arena_gb']:.0f} GB of RAM{fit}")
|
||||
rec = str(names.index("IQ3_XXS") + 1) if ram >= 60 and "IQ3_XXS" in names else "1"
|
||||
model = a.model or names[int(ask("Which size?", [str(i) for i in range(1, len(names) + 1)], rec, a.yes)) - 1]
|
||||
budget = None
|
||||
budget, q4_split = None, False
|
||||
if MODELS[model].get("budget"):
|
||||
# Unsloth's UD-Q4_K_XL: a RAM budget of experts, the rest from the GGUF on the SSD - not the low-RAM mode (no
|
||||
# experts.bin: it would be another 77 GB on the disk), and one GPU (the budget mode has no layer split)
|
||||
# experts.bin: it would be another 77 GB on the disk), and one GPU (the budget mode has no layer split) unless
|
||||
# the RAM holds the GGUFs and 24 GB more: then several, without the budget, if asked for (#498)
|
||||
warn(f"{model} is EXPERIMENTAL (docs/UNSLOTH_Q4.md): most of its experts are read from the SSD while it "
|
||||
"answers, so it is several times slower than the 2-3-bit models; quality checked against llama.cpp")
|
||||
if hip:
|
||||
@@ -2831,12 +3158,13 @@ def main() -> int:
|
||||
"install it anyway", " Install it anyway?")
|
||||
warn(f"installing {model} with {ram:.0f} GB of RAM, as you chose")
|
||||
budget = budget_choice(model, ram, a.resident_budget_gib)
|
||||
ok(f"RAM budget: {budget:g} GiB of {model}'s experts in RAM, the rest read from the model files on the SSD")
|
||||
if multi and not unsloth_together(a, model, ram, gpu, chosen):
|
||||
multi, sel, chosen = [], [gpu["index"]], [gpu]
|
||||
q4_split = bool(multi) # #498: on several GPUs without the RAM budget
|
||||
if not q4_split:
|
||||
ok(f"RAM budget: {budget:g} GiB of {model}'s experts in RAM, the rest read from the model files on the SSD")
|
||||
if a.low_ram not in ("auto", "off"):
|
||||
warn(f"--low-ram {a.low_ram} does not apply to {model}: it always reads part of its experts from the files")
|
||||
if multi:
|
||||
warn(f"{model} runs on one GPU: using " + gpu_name(gpu) + " only")
|
||||
multi, sel, chosen = [], [gpu["index"]], [gpu]
|
||||
elif a.resident_budget_gib is not None:
|
||||
warn(f"--resident-budget-gib is for UD-Q4_K_XL: {model} keeps all of its experts in RAM or in the low-RAM mode")
|
||||
low_ram = budget is None and (a.low_ram in ("on", "resident", "mmap") or
|
||||
@@ -2975,6 +3303,9 @@ def main() -> int:
|
||||
models_dir = Path(a.gguf_dir) if a.gguf_dir else Path(a.models_dir) / tag
|
||||
shards = gguf_dir_shards(models_dir, fam, model) if a.gguf_dir else \
|
||||
[models_dir / fam["file"].format(q=model, i=i) for i in range(1, fam.get("shards", 2) + 1)]
|
||||
problem = gguf_dir_problem(models_dir, shards[0], fam, model) if a.gguf_dir else None
|
||||
if problem: # #444: files Strata cannot run, or another choice's files
|
||||
fail(*problem)
|
||||
if not a.gguf_dir and not all(sh.exists() and done(sh) for sh in shards):
|
||||
for r in elsewhere: # already downloaded in a Strata folder on another drive
|
||||
cand = [r / "models" / tag / sh.name for sh in shards]
|
||||
@@ -3040,6 +3371,13 @@ def main() -> int:
|
||||
# ---- 5. the model files
|
||||
step(5, f"downloading {fam['title']} {model}")
|
||||
if not a.gguf_dir:
|
||||
missing = [s.name for s in shards if not (s.exists() and done(s))]
|
||||
if missing: # #495: files downloaded by hand go here, or --gguf-dir
|
||||
say(f" The model files go in {models_dir}")
|
||||
say(f" Files you already have: put them here with their original names ({', '.join(missing)}), or use "
|
||||
"--gguf-dir <their folder>.")
|
||||
if hf_endpoint() != HF_DEFAULT:
|
||||
say(f" Downloading from {hf_endpoint()} (HF_ENDPOINT)")
|
||||
for s in shards:
|
||||
if s.exists() and done(s):
|
||||
ok(f"{s.name} already downloaded")
|
||||
@@ -3112,6 +3450,8 @@ def main() -> int:
|
||||
draft_vocab = a.draft_vocab or saved_draft_vocab(ROOT / f"strata-{tag.lower()}.json")
|
||||
refresh_draft_vocab(rt, draft_vocab or "cjk")
|
||||
ok(f"MTP draft layer: {rt}")
|
||||
for line in draft_vocab_note(gpu.get("vram_gb", 0.0), draft_vocab): # #474: a recommendation, nothing changes
|
||||
say(" " + line)
|
||||
|
||||
# ---- 7. the start script
|
||||
step(7, "writing the start script")
|
||||
@@ -3163,7 +3503,9 @@ def main() -> int:
|
||||
warn(f"KV streaming needs ~{kv_ram_gb:.1f} GB of RAM beside the ~{MODELS[model]['ram_gb']} GB {model} "
|
||||
f"uses; this PC has {ram:.0f}. Kept as you chose (--kv-streaming on): it may page or run out of RAM "
|
||||
"under load")
|
||||
if budget is not None and a.resident_budget_gib is None: # its RAM comes out of the experts' budget
|
||||
if q4_split: # #498: no budget to take it out of
|
||||
pass
|
||||
elif budget is not None and a.resident_budget_gib is None: # its RAM comes out of the experts' budget
|
||||
budget = resident_budget_gib(model, ram, kv_ram_gb)
|
||||
ok(f"RAM budget: {budget} GiB (less the KV cache's RAM)")
|
||||
elif budget is not None and budget > resident_budget_gib(model, ram, kv_ram_gb):
|
||||
@@ -3171,10 +3513,21 @@ def main() -> int:
|
||||
f"would take them out of it: {resident_budget_gib(model, ram, kv_ram_gb)} GiB); kept as you chose")
|
||||
elif a.kv_streaming == "on":
|
||||
warn("--kv-streaming on: a context under 64K is not streamed (the attention's window holds all of it): off")
|
||||
if budget is not None: # UD-Q4_K_XL: the experts read from the GGUF in place, the most-used N GiB kept in RAM
|
||||
args += ["--resident-budget-gib", f"{budget:g}"]
|
||||
if budget is not None and not q4_split: # UD-Q4_K_XL: the experts read from the GGUF in place, the most-used N
|
||||
args += ["--resident-budget-gib", f"{budget:g}"] # GiB kept in RAM (#498: a layer split has no budget)
|
||||
if vision != "none":
|
||||
args += ["--vision", "--vram-reserve-mib", str(VISION[vision]["reserve_mib"])]
|
||||
if a.vram_reserve_mib is not None: # #493: VRAM left free for other programs (only when given)
|
||||
if "--vram-reserve-mib" in args:
|
||||
i = args.index("--vram-reserve-mib") + 1
|
||||
if vision == "gpu" and a.vram_reserve_mib < int(args[i]):
|
||||
warn(f"--vram-reserve-mib {a.vram_reserve_mib}: the image encoder on the GPU needs ~{args[i]} MiB of "
|
||||
"it; kept as you chose (it may run out of VRAM when it reads a picture)")
|
||||
args[i] = str(a.vram_reserve_mib)
|
||||
else:
|
||||
args += ["--vram-reserve-mib", str(a.vram_reserve_mib)]
|
||||
ok(f"VRAM kept free for other programs: {a.vram_reserve_mib} MiB (--vram-reserve-mib; the expert cache takes "
|
||||
"that much less)")
|
||||
if esp is not None:
|
||||
# the package's profile, with llama.cpp's flags (the engine takes the same ones)
|
||||
args += ["--control-vector-scaled", f"{esp}:1.0", "--control-vector-layer-range", "4", "44",
|
||||
@@ -3217,13 +3570,15 @@ def main() -> int:
|
||||
import calibrate as CAL
|
||||
cfg["args"] = CAL.apply(cfg["args"], cal.get("settings") or {})
|
||||
ok("the settings tuned for this PC earlier are used" + (f" ({cal['date']})" if cal.get("date") else ""))
|
||||
cfg_path.write_text(json.dumps(cfg, indent=1), encoding="utf-8")
|
||||
write_config(cfg_path, cfg)
|
||||
script = write_run_script(tag, cfg_path, port)
|
||||
# offered only when someone answers: --yes installs and adopted earlier installs are not held up by it
|
||||
if cal is None and not hip and not a.no_start and not a.yes and ask(
|
||||
"Tune Strata for this PC now? It measures a few engine settings (about 5-10 minutes; the PC is busy "
|
||||
"meanwhile; later: START-HERE --calibrate)", ["y", "n"], "y", a.yes) == "y":
|
||||
calibrate_config(cfg_path)
|
||||
tuned = calibrate_config(cfg_path)
|
||||
else:
|
||||
tuned = None # not asked for: nothing to repeat below
|
||||
ok(f"start script: {script.name}")
|
||||
|
||||
say()
|
||||
@@ -3236,6 +3591,9 @@ def main() -> int:
|
||||
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 tuned is False: # #447: a failed tuning is repeated here, not only above
|
||||
say(" Tuning: FAILED (the reason is above): the default settings stay - "
|
||||
f"{'START-HERE.bat' if WIN else './setup.sh'} --calibrate tries again")
|
||||
if a.no_start:
|
||||
return 0
|
||||
return start(cfg_path, port)
|
||||
|
||||
@@ -101,6 +101,97 @@ bool device_summary(int ordinal, std::string& name, std::string& detail) {
|
||||
return true;
|
||||
}
|
||||
|
||||
} // namespace strata::core
|
||||
|
||||
#if defined(STRATA_USE_HIP) && defined(_WIN32)
|
||||
// The free VRAM figure on Windows HIP (include/strata/hip_compat/cuda_runtime.h maps cudaMemGetInfo here).
|
||||
//
|
||||
// hipMemGetInfo (ROCclr's PAL backend, Device::globalFreeMemory) is the card's size minus this process's own
|
||||
// allocations: it asks Windows for this process's usage only, never for what the desktop and other programs hold.
|
||||
// So on a card that also drives the desktop the engine counted ~930 MiB that was not there, the expert cache filled
|
||||
// the card, and WDDM moved memory out to system RAM: decode at 30 tok/s instead of 41 (RX 6800, HIP SDK 7.2).
|
||||
// Windows itself gives each process a video memory budget (DXGI QueryVideoMemoryInfo) that does account for the
|
||||
// others. Measured on that card: the budget sits 0.8 GiB below the card's size while this process is small and
|
||||
// 1.8 GiB below once it holds 11 GiB; allocations past it still succeed, and are what Windows moves out. The
|
||||
// free figure here is hipMemGetInfo's lowered by what the budget withholds (the card's size minus the budget).
|
||||
// Subtracting from HIP's own figure, rather than taking the budget minus DXGI's usage, keeps memory the HIP runtime
|
||||
// has freed and holds in its cache counted as free, as hipMemGetInfo counts it. dxgi.dll is loaded at run time, so
|
||||
// nothing new is linked; when it or the card's adapter cannot be found, hipMemGetInfo's figure stands.
|
||||
#ifndef WIN32_LEAN_AND_MEAN
|
||||
#define WIN32_LEAN_AND_MEAN
|
||||
#endif
|
||||
#ifndef NOMINMAX
|
||||
#define NOMINMAX
|
||||
#endif
|
||||
#include <windows.h>
|
||||
#include <dxgi1_4.h>
|
||||
#include <atomic>
|
||||
#include <cstdlib>
|
||||
#include <mutex>
|
||||
|
||||
namespace strata::hip_compat {
|
||||
namespace {
|
||||
IDXGIAdapter3* budget_adapter(int device) {
|
||||
constexpr int kMaxDevices = 16;
|
||||
static std::mutex mu;
|
||||
static IDXGIAdapter3* adapters[kMaxDevices] = {};
|
||||
static bool tried[kMaxDevices] = {};
|
||||
if (device < 0 || device >= kMaxDevices) return nullptr;
|
||||
std::lock_guard<std::mutex> lock(mu);
|
||||
if (tried[device]) return adapters[device];
|
||||
tried[device] = true;
|
||||
hipDeviceProp_t p{};
|
||||
if (hipGetDeviceProperties(&p, device) != hipSuccess) {
|
||||
(void) hipGetLastError();
|
||||
return nullptr;
|
||||
}
|
||||
LUID luid{};
|
||||
std::memcpy(&luid, p.luid, sizeof luid);
|
||||
using CreateFactory = HRESULT(WINAPI*)(REFIID, void**);
|
||||
static HMODULE dxgi = LoadLibraryW(L"dxgi.dll"); // kept for the process' lifetime, like the adapters
|
||||
const auto create = dxgi ? (CreateFactory) (void*) GetProcAddress(dxgi, "CreateDXGIFactory1") : nullptr;
|
||||
IDXGIFactory4* factory = nullptr;
|
||||
if (create == nullptr || FAILED(create(__uuidof(IDXGIFactory4), (void**) &factory))) return nullptr;
|
||||
IDXGIAdapter3* adapter = nullptr;
|
||||
if (FAILED(factory->EnumAdapterByLuid(luid, __uuidof(IDXGIAdapter3), (void**) &adapter))) adapter = nullptr;
|
||||
factory->Release();
|
||||
adapters[device] = adapter;
|
||||
return adapter;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
hipError_t mem_get_info(size_t* free_bytes, size_t* total_bytes) {
|
||||
const hipError_t e = hipMemGetInfo(free_bytes, total_bytes);
|
||||
static const bool off = [] {
|
||||
const char* v = std::getenv("STRATA_WDDM_BUDGET");
|
||||
return v != nullptr && std::atoi(v) == 0;
|
||||
}();
|
||||
if (e != hipSuccess || off || free_bytes == nullptr || total_bytes == nullptr) return e;
|
||||
int device = 0;
|
||||
if (hipGetDevice(&device) != hipSuccess) {
|
||||
(void) hipGetLastError();
|
||||
return e;
|
||||
}
|
||||
IDXGIAdapter3* adapter = budget_adapter(device);
|
||||
DXGI_QUERY_VIDEO_MEMORY_INFO local{};
|
||||
if (adapter == nullptr || FAILED(adapter->QueryVideoMemoryInfo(0, DXGI_MEMORY_SEGMENT_GROUP_LOCAL, &local)) ||
|
||||
local.Budget == 0 || local.Budget >= *total_bytes)
|
||||
return e;
|
||||
const size_t withheld = *total_bytes - (size_t) local.Budget; // the desktop's and other programs' share
|
||||
static std::atomic<bool> said{false};
|
||||
if (!said.exchange(true)) {
|
||||
std::fprintf(stderr, "strata: Windows budgets %llu of this card's %llu MiB for this process; free VRAM is "
|
||||
"counted within that (STRATA_WDDM_BUDGET=0: off)\n",
|
||||
(unsigned long long) (local.Budget >> 20), (unsigned long long) (*total_bytes >> 20));
|
||||
}
|
||||
*free_bytes = *free_bytes > withheld ? *free_bytes - withheld : 0;
|
||||
return e;
|
||||
}
|
||||
} // namespace strata::hip_compat
|
||||
#endif
|
||||
|
||||
namespace strata::core {
|
||||
|
||||
std::string gpu_arch_problem(int ordinal) {
|
||||
#if defined(STRATA_USE_HIP)
|
||||
int count = 0;
|
||||
|
||||
@@ -9,6 +9,12 @@
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#if defined(STRATA_USE_HIP) && defined(_WIN32)
|
||||
#ifndef NOMINMAX
|
||||
#define NOMINMAX
|
||||
#endif
|
||||
#include <windows.h>
|
||||
#endif
|
||||
|
||||
static std::string human(uint64_t b) {
|
||||
char buf[64];
|
||||
@@ -68,6 +74,13 @@ int main(int argc, char** argv) {
|
||||
std::printf(" multiprocessors %d\n", d.multi_processor_count);
|
||||
std::printf(" VRAM total / free %s / %s\n", human(d.total_bytes).c_str(), human(d.free_bytes).c_str());
|
||||
std::printf(" driver / runtime %d / %d\n", d.driver_version, d.runtime_version);
|
||||
#if defined(STRATA_USE_HIP) && defined(_WIN32)
|
||||
// #468 #461: which HIP runtime this process loaded - the one beside the exe, or an AMD driver's System32 copy
|
||||
if (HMODULE h = GetModuleHandleA("amdhip64_7.dll")) {
|
||||
char path[MAX_PATH] = {};
|
||||
if (GetModuleFileNameA(h, path, MAX_PATH) > 0) std::printf(" HIP runtime %s\n", path);
|
||||
}
|
||||
#endif
|
||||
|
||||
// The planner's view against the card's. A plan that does not fit in what is actually FREE is the
|
||||
// failure this print exists to make visible at startup rather than at token 4000.
|
||||
|
||||
@@ -3,7 +3,9 @@
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdio>
|
||||
#include <filesystem>
|
||||
#include <utility>
|
||||
#include <cstring>
|
||||
|
||||
@@ -67,6 +69,77 @@ bool read_expert_profile(const std::string& path, int64_t n_layers, int64_t n_ex
|
||||
return true;
|
||||
}
|
||||
|
||||
std::vector<std::pair<int32_t, int32_t>> rank_learned_profile(int64_t n_layers, int64_t n_expert,
|
||||
const std::vector<uint8_t>& resident,
|
||||
const std::vector<double>& heat,
|
||||
const std::vector<std::pair<int32_t, int32_t>>& prior) {
|
||||
const size_t n = (size_t) (n_layers * n_expert);
|
||||
std::vector<int64_t> prior_rank(n, INT64_MAX);
|
||||
for (size_t r = 0; r < prior.size(); ++r) {
|
||||
const auto [l, e] = prior[r];
|
||||
if (l >= 0 && l < n_layers && e >= 0 && e < n_expert) {
|
||||
int64_t& pr = prior_rank[(size_t) (l * n_expert + e)];
|
||||
if (pr == INT64_MAX) pr = (int64_t) r;
|
||||
}
|
||||
}
|
||||
std::vector<int64_t> order(n);
|
||||
for (size_t i = 0; i < n; ++i) order[i] = (int64_t) i;
|
||||
auto res = [&](int64_t i) { return (size_t) i < resident.size() && resident[(size_t) i] != 0; };
|
||||
auto ht = [&](int64_t i) { return (size_t) i < heat.size() ? heat[(size_t) i] : 0.0; };
|
||||
std::stable_sort(order.begin(), order.end(), [&](int64_t a, int64_t b) {
|
||||
if (res(a) != res(b)) return res(a);
|
||||
if (ht(a) != ht(b)) return ht(a) > ht(b);
|
||||
if (prior_rank[(size_t) a] != prior_rank[(size_t) b]) return prior_rank[(size_t) a] < prior_rank[(size_t) b];
|
||||
return a < b;
|
||||
});
|
||||
std::vector<std::pair<int32_t, int32_t>> ranked(n);
|
||||
for (size_t r = 0; r < n; ++r)
|
||||
ranked[r] = {(int32_t) (order[r] / n_expert), (int32_t) (order[r] % n_expert)};
|
||||
return ranked;
|
||||
}
|
||||
|
||||
bool write_expert_profile(const std::string& path, int64_t n_layers, int64_t n_expert,
|
||||
const std::vector<std::pair<int32_t, int32_t>>& ranked, std::string& err) {
|
||||
if (n_layers <= 0 || n_expert <= 0 || n_layers > 65535 || n_expert > 65535) {
|
||||
err = "write_expert_profile: the model's layout does not fit the format";
|
||||
return false;
|
||||
}
|
||||
std::vector<int32_t> table((size_t) (n_layers * n_expert), -1);
|
||||
std::vector<uint16_t> pairs;
|
||||
pairs.reserve(ranked.size() * 2);
|
||||
for (size_t r = 0; r < ranked.size(); ++r) {
|
||||
const auto [l, e] = ranked[r];
|
||||
if (l < 0 || l >= n_layers || e < 0 || e >= n_expert) {
|
||||
err = "write_expert_profile: a ranked pair is out of range";
|
||||
return false;
|
||||
}
|
||||
table[(size_t) (l * n_expert + e)] = (int32_t) r;
|
||||
pairs.push_back((uint16_t) l);
|
||||
pairs.push_back((uint16_t) e);
|
||||
}
|
||||
const uint32_t hdr[5] = {1u, (uint32_t) n_layers, (uint32_t) n_expert, (uint32_t) ranked.size(),
|
||||
(uint32_t) ranked.size()};
|
||||
const std::string tmp = path + ".tmp";
|
||||
std::FILE* f = std::fopen(tmp.c_str(), "wb");
|
||||
if (f == nullptr) {
|
||||
err = "write_expert_profile: cannot create " + tmp;
|
||||
return false;
|
||||
}
|
||||
// the format is little-endian (make_profile.py's "<"): so is every machine this engine runs on
|
||||
bool ok = std::fwrite("STRP", 1, 4, f) == 4 && std::fwrite(hdr, 4, 5, f) == 5 &&
|
||||
(pairs.empty() || std::fwrite(pairs.data(), 2, pairs.size(), f) == pairs.size()) &&
|
||||
std::fwrite(table.data(), 4, table.size(), f) == table.size();
|
||||
ok = (std::fclose(f) == 0) && ok;
|
||||
std::error_code ec;
|
||||
if (ok) std::filesystem::rename(tmp, path, ec); // replaces an existing file (MoveFileEx / rename(2))
|
||||
if (!ok || ec) {
|
||||
std::filesystem::remove(tmp, ec);
|
||||
err = "write_expert_profile: cannot write " + path;
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
ExpertCache::~ExpertCache() { close(); }
|
||||
|
||||
#if defined(STRATA_USE_HIP)
|
||||
|
||||
@@ -1024,6 +1024,22 @@ bool FileExpertSource::pin_cache_complement(
|
||||
uint64_t bytes = 0;
|
||||
if (!detail::make_cache_complement_plan(n_layers_, n_expert_, layer_blob_bytes_, primary_gpu_pairs,
|
||||
additional_gpu_pairs, offsets, bytes, err)) return false;
|
||||
#if defined(_WIN32)
|
||||
// #467: the GPU cache's pre-fill touched its experts through the mapping (~19 GiB on a 24 GB card), and Windows
|
||||
// counts those file pages in this process's working set, not as available: a 32 GB PC read 0.44 GiB here
|
||||
// (20.7 GiB before the start). Trimmed, they move to the standby list (still cached, counted as available).
|
||||
// Resident mode only: nothing else calls this function. Locked/pinned pages stay; the rest fault back softly.
|
||||
{
|
||||
uint64_t before = 0, after = 0;
|
||||
const bool read_before = available_memory_bytes(before);
|
||||
(void) SetProcessWorkingSetSize(GetCurrentProcess(), (SIZE_T) -1, (SIZE_T) -1);
|
||||
if (read_before && available_memory_bytes(after))
|
||||
std::fprintf(stderr, "FileExpertSource: available RAM %.2f GiB, %.2f GiB after the mapped experts left the "
|
||||
"process working set (#467)\n",
|
||||
(double) before / 1073741824.0, (double) after / 1073741824.0);
|
||||
}
|
||||
#endif
|
||||
const bool what_fits = budget_bytes == kResidentWhatFits; // #467: the soft mode's second try
|
||||
uint64_t budget_physical = 0; // #403: the RAM reading a budget was sized from (0: no budget)
|
||||
if (budget_bytes > 0) {
|
||||
// CS-T: a RAM budget. The complement's experts in `rank` order (the expert profile, hottest first) while
|
||||
@@ -1041,10 +1057,16 @@ bool FileExpertSource::pin_cache_complement(
|
||||
// lower). An unclamped budget is unchanged.
|
||||
const uint64_t margin = 256ull << 20;
|
||||
const uint64_t clamped = room > margin ? room - margin : 0;
|
||||
std::fprintf(stderr, "FileExpertSource: --resident-budget-gib %.2f is more than the RAM has room for "
|
||||
"(%.2f GiB available minus %.0f GiB headroom and a 0.25 GiB margin): %.2f GiB\n",
|
||||
(double) budget_bytes / 1073741824.0, (double) physical / 1073741824.0,
|
||||
(double) headroom_bytes / 1073741824.0, (double) clamped / 1073741824.0);
|
||||
if (what_fits)
|
||||
std::fprintf(stderr, "FileExpertSource: RAM room for the complement: %.2f GiB (%.2f GiB available "
|
||||
"minus %.0f GiB headroom and a 0.25 GiB margin)\n",
|
||||
(double) clamped / 1073741824.0, (double) physical / 1073741824.0,
|
||||
(double) headroom_bytes / 1073741824.0);
|
||||
else
|
||||
std::fprintf(stderr, "FileExpertSource: --resident-budget-gib %.2f is more than the RAM has room for "
|
||||
"(%.2f GiB available minus %.0f GiB headroom and a 0.25 GiB margin): %.2f GiB\n",
|
||||
(double) budget_bytes / 1073741824.0, (double) physical / 1073741824.0,
|
||||
(double) headroom_bytes / 1073741824.0, (double) clamped / 1073741824.0);
|
||||
budget_bytes = clamped;
|
||||
}
|
||||
std::vector<uint64_t> ranked(offsets.size(), kNoComplement);
|
||||
@@ -1061,6 +1083,10 @@ bool FileExpertSource::pin_cache_complement(
|
||||
at += b;
|
||||
++held;
|
||||
}
|
||||
if (what_fits && held == 0) { // #467: nothing to keep - the caller's plain mmap fallback, not an empty copy
|
||||
err = "FileExpertSource: the RAM has no room for any expert of the complement";
|
||||
return false;
|
||||
}
|
||||
std::fprintf(stderr, "FileExpertSource: RAM budget %.2f GiB: %lld of the %.2f GiB of experts the GPU cache does "
|
||||
"not hold, by profile rank; the rest are read from the files\n",
|
||||
(double) budget_bytes / 1073741824.0, (long long) held, (double) bytes / 1073741824.0);
|
||||
|
||||
@@ -372,6 +372,45 @@ void MtpDrafter::set_draft_history(const int32_t* tail, int64_t n_tail, int32_t
|
||||
coupled_hist_base(tail, n_tail, next, h, h_chist_ + (kCoupledHistCap - h));
|
||||
}
|
||||
|
||||
namespace {
|
||||
// #474: what to change when the draft head's token subset does not fit the VRAM left - a smaller subset that setup
|
||||
// ships (data/draft_vocab_*.bin; their token counts below), and how to pick it. The start used to stop at "the draft
|
||||
// head does not fit" with no hint. Text only, after the failure: nothing changes for a start that fits.
|
||||
void draft_head_hint(int64_t n_tokens, int64_t row_bytes) {
|
||||
size_t free_b = 0, total_b = 0;
|
||||
const bool have_free = cudaMemGetInfo(&free_b, &total_b) == cudaSuccess;
|
||||
auto mib = [&](int64_t n) { return (double) (n * row_bytes) / 1048576.0; };
|
||||
std::string free_s;
|
||||
if (have_free) {
|
||||
char b[64];
|
||||
std::snprintf(b, sizeof(b), " and %.0f MiB is free", (double) free_b / 1048576.0);
|
||||
free_s = b;
|
||||
}
|
||||
std::fprintf(stderr, "strata mtp: the draft head over %lld tokens needs %.0f MiB of VRAM%s.\n",
|
||||
(long long) n_tokens, mib(n_tokens), free_s.c_str());
|
||||
struct Subset { const char* name; int64_t tokens; const char* what; };
|
||||
static const Subset smaller[] = {{"cyrillic", 58963, "English, code and the Cyrillic script"},
|
||||
{"en", 40525, "English and code"}};
|
||||
std::string opts;
|
||||
for (const Subset& s : smaller)
|
||||
if (s.tokens < n_tokens) {
|
||||
char b[160];
|
||||
std::snprintf(b, sizeof(b), "%s--draft-vocab %s (%s, ~%.0f MiB)", opts.empty() ? "" : " or ", s.name,
|
||||
s.what, mib(s.tokens));
|
||||
opts += b;
|
||||
}
|
||||
if (!opts.empty())
|
||||
std::fprintf(stderr, "strata mtp: hint: a smaller draft vocabulary needs less VRAM: %s. Start once with it - "
|
||||
"START-HERE.bat --draft-vocab en (Windows) or ./setup.sh --draft-vocab en - and the model "
|
||||
"keeps it (\"draft_vocab\" in its strata-*.json config); or a smaller --context in "
|
||||
"setup.\n",
|
||||
opts.c_str());
|
||||
else
|
||||
std::fprintf(stderr, "strata mtp: hint: this is already the smallest shipped draft vocabulary: a smaller "
|
||||
"--context (or closing what else uses the GPU) leaves it room.\n");
|
||||
}
|
||||
} // namespace
|
||||
|
||||
bool MtpDrafter::bind(const WeightTable& wt, const NativeHead* head, const float* window_R, std::string& err) {
|
||||
const OnDevice on_device(device_);
|
||||
wt_ = &wt;
|
||||
@@ -395,6 +434,7 @@ bool MtpDrafter::bind(const WeightTable& wt, const NativeHead* head, const float
|
||||
if (cudaMalloc((void**) &dvocab_, raw.size()) != cudaSuccess ||
|
||||
cudaMalloc((void**) &dhead_, (size_t) (n_dvocab_ * row_bytes)) != cudaSuccess) {
|
||||
err = "mtp: the draft head does not fit";
|
||||
draft_head_hint(n_dvocab_, row_bytes);
|
||||
return false;
|
||||
}
|
||||
cudaMemcpy(dvocab_, raw.data(), raw.size(), cudaMemcpyHostToDevice);
|
||||
|
||||
@@ -644,6 +644,229 @@ __global__ void __launch_bounds__(SCORE_WARPS * 32) block_scores_multi_kernel(co
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if !defined(__HIPCC__)
|
||||
// ---- the decode top-k on a thread-block cluster (sm_90+; S19). One CTA per query (block_topk_reg_kernel, or
|
||||
// block_topk_kernel above 33,792 blocks: a --max-context over ~135K) makes four radix passes and two scans over up to
|
||||
// 65,538 blocks on ONE SM while the rest of the GPU idles - a decode window has 1-5 queries. Here a cluster of CL_N
|
||||
// CTAs shares a query: CTA r holds the keys of blocks [r * per, (r + 1) * per) in shared memory, builds the digit
|
||||
// histogram of its keys, and PUSHES it into slot r of every CTA's `hin` (distributed shared memory); after one cluster
|
||||
// barrier each CTA sums the CL_N slots itself, so all of them take the same digit. The emit pass needs each CTA's
|
||||
// cells above / at the threshold: pushed the same way, then each CTA offsets its own by the ranks before it.
|
||||
// RTX 5070, per call at capacity = context (decode_cluster_parity --bench): 21.9 -> 15.6 us at 32K, 58 -> 18 at 128K,
|
||||
// 200 -> 22 at 262K (the one-CTA kernels' digit search was serial, too: one thread over 256 bins per pass).
|
||||
//
|
||||
// IDENTICAL IDS, by construction rather than by luck: the threshold thr (the width-th largest key, cells counted with
|
||||
// their weights) and `above` (cells with a larger key) are pure functions of the multiset of (key, weight) - integer
|
||||
// histograms, summed in any order - and the digit rule is block_topk_kernel's, written as a scan (the largest digit d
|
||||
// whose cells at or above it reach `need`; 0 when none). The cells emitted are then fixed: every cell of a block with
|
||||
// key > thr, and the first eq_budget = width - above cells at thr in ascending order (a block at thr may be cut),
|
||||
// written in ascending cell order. Position of a thread's first cell = (cells above thr before it) + min(cells at thr
|
||||
// before it, eq_budget): the telescoped sum of the reference's per-thread clamp. NaN keys are 0 (order_key) as there.
|
||||
//
|
||||
// Barriers: barrier.cluster arrive / wait (each CTA's threads all take part; it also orders the CTA's own shared
|
||||
// memory, so it doubles as __syncthreads). Phase 0 is a relaxed arrive at entry, waited before the first remote store
|
||||
// (a CTA's shared memory exists once it runs). Pushes go out before the arrive that releases them, and nothing reads
|
||||
// another CTA's memory after the last wait - so a CTA may exit without a closing barrier (a wait counts the threads
|
||||
// that have not exited). `hin` alternates by pass parity: CTA x reads slot buffer p&1 in pass p before arriving at
|
||||
// pass p+1's barrier, and nobody writes buffer p&1 again (pass p+2) before waiting on that barrier.
|
||||
constexpr int CL_N = 8; // CTAs per query (the portable cluster size)
|
||||
constexpr int CL_T = 1024; // threads per CTA
|
||||
constexpr int CL_MAXQ = 16; // larger calls (prefill sub-batches) fill the GPU with one CTA per query already
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
#define STRATA_SEL_CLUSTER 1
|
||||
#else
|
||||
#define STRATA_SEL_CLUSTER 0 // an older target's code is a trap; qsa_block_topk_cluster never launches it there
|
||||
#endif
|
||||
|
||||
#if STRATA_SEL_CLUSTER
|
||||
__device__ __forceinline__ void cl_arrive_relaxed() {
|
||||
asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" ::: "memory");
|
||||
}
|
||||
__device__ __forceinline__ void cl_arrive() { asm volatile("barrier.cluster.arrive.release.aligned;\n" ::: "memory"); }
|
||||
__device__ __forceinline__ void cl_wait() { asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); }
|
||||
__device__ __forceinline__ unsigned cl_rank() {
|
||||
unsigned r;
|
||||
asm volatile("mov.u32 %0, %%cluster_ctarank;\n" : "=r"(r));
|
||||
return r;
|
||||
}
|
||||
template <typename T> __device__ __forceinline__ T* cl_map(T* p, unsigned rank) { // p in CTA `rank`'s shared memory
|
||||
uint64_t o;
|
||||
asm volatile("mapa.u64 %0, %1, %2;\n" : "=l"(o) : "l"((uint64_t) p), "r"(rank));
|
||||
return reinterpret_cast<T*>(o);
|
||||
}
|
||||
// exclusive block scan of a 64-bit value over CL_T threads; `total` = the CTA's sum
|
||||
__device__ __forceinline__ unsigned long long cl_excl_scan64(unsigned long long v, unsigned long long* s_warp,
|
||||
unsigned long long& total) {
|
||||
const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
|
||||
unsigned long long x = v;
|
||||
#pragma unroll
|
||||
for (int o = 1; o < 32; o <<= 1) {
|
||||
const unsigned long long y = __shfl_up_sync(0xffffffffu, x, o);
|
||||
if (lane >= o) x += y;
|
||||
}
|
||||
if (lane == 31) s_warp[warp] = x;
|
||||
__syncthreads();
|
||||
if (warp == 0) {
|
||||
const unsigned long long w = s_warp[lane];
|
||||
unsigned long long z = w;
|
||||
#pragma unroll
|
||||
for (int o = 1; o < 32; o <<= 1) {
|
||||
const unsigned long long y = __shfl_up_sync(0xffffffffu, z, o);
|
||||
if (lane >= o) z += y;
|
||||
}
|
||||
s_warp[lane] = z - w;
|
||||
if (lane == 31) s_warp[32] = z;
|
||||
}
|
||||
__syncthreads();
|
||||
total = s_warp[32];
|
||||
return s_warp[warp] + x - v;
|
||||
}
|
||||
#endif
|
||||
|
||||
// grid (CL_N, nq), cluster (CL_N, 1, 1), CL_T threads, dynamic shared memory: ceil(max_blocks / CL_N) keys
|
||||
__global__ void __launch_bounds__(CL_T) block_topk_cluster_kernel(const float* __restrict__ scores,
|
||||
const int32_t* __restrict__ steps, int64_t max_blocks,
|
||||
int64_t cap, int32_t* __restrict__ ids) {
|
||||
#if STRATA_SEL_CLUSTER
|
||||
extern __shared__ uint32_t keys[]; // this CTA's blocks' keys
|
||||
__shared__ __align__(16) int hin[2][CL_N][256]; // the cluster's histograms, slot = the pushing rank
|
||||
__shared__ __align__(16) int hloc[2][256]; // this CTA's histogram
|
||||
__shared__ unsigned long long cnt[CL_N]; // per rank: its cells above thr (low 32 bits) and at it (high)
|
||||
__shared__ unsigned long long s_warp[33];
|
||||
__shared__ uint32_t s_prefix;
|
||||
__shared__ int s_above;
|
||||
const unsigned rank = cl_rank();
|
||||
const int t = threadIdx.x, lane = t & 31, warp = t >> 5;
|
||||
cl_arrive_relaxed(); // phase 0: this CTA runs
|
||||
const int64_t qi = blockIdx.y;
|
||||
const int32_t* st = steps + qi * kStepCount;
|
||||
const int64_t n_kv = st[kStepNKv], n_bid = st[kStepNBid], width = st[kStepWidth];
|
||||
int32_t* out = ids + qi * cap;
|
||||
if (n_kv <= width) { // the identity, split over the cluster; no remote access
|
||||
for (int64_t j = (int64_t) rank * CL_T + t; j < n_kv; j += (int64_t) CL_N * CL_T) out[j] = (int32_t) j;
|
||||
return;
|
||||
}
|
||||
// the tail block n_bid weighs its cells (0..3); with none it is not a candidate (the reference skips w == 0)
|
||||
const int wtail = (int) (n_kv - n_bid * R);
|
||||
const int64_t nbe = n_bid + (wtail > 0 ? 1 : 0);
|
||||
const int64_t per = (nbe + CL_N - 1) / CL_N; // <= ceil(max_blocks / CL_N), the keys the host sized
|
||||
const int64_t lo = (int64_t) rank * per < nbe ? (int64_t) rank * per : nbe;
|
||||
const int64_t hi = lo + per < nbe ? lo + per : nbe;
|
||||
const int n = (int) (hi - lo);
|
||||
const int tail_i = (wtail > 0 && n_bid >= lo && n_bid < hi) ? (int) (n_bid - lo) : -1;
|
||||
const float* sc = scores + qi * max_blocks + lo;
|
||||
for (int i = t; i < n; i += CL_T) keys[i] = order_key(sc[i]);
|
||||
for (int i = t; i < 2 * 256; i += CL_T) (&hloc[0][0])[i] = 0;
|
||||
__syncthreads();
|
||||
uint32_t prefix = 0;
|
||||
int above = 0; // cells strictly above the digits fixed so far
|
||||
for (int pass = 0; pass < 4; ++pass) {
|
||||
const int shift = 24 - 8 * pass;
|
||||
const uint32_t hmask = pass == 0 ? 0u : (0xffffffffu << (shift + 8));
|
||||
int* hl = hloc[pass & 1];
|
||||
// warp-aggregated: the scores share their top bits, so plain atomics would queue on a few bins
|
||||
for (int i0 = 0; i0 < n; i0 += CL_T) {
|
||||
const int i = i0 + t;
|
||||
int bin = -1;
|
||||
if (i < n) {
|
||||
const uint32_t k = keys[i];
|
||||
if ((k & hmask) == prefix) bin = (int) ((k >> shift) & 255);
|
||||
}
|
||||
const unsigned same = __match_any_sync(0xffffffffu, bin);
|
||||
if (bin >= 0 && lane == __ffs(same) - 1) atomicAdd(&hl[bin], __popc(same) * R);
|
||||
if (bin >= 0 && i == tail_i) atomicAdd(&hl[bin], wtail - R);
|
||||
}
|
||||
__syncthreads();
|
||||
if (pass == 0) cl_wait(); // phase 0 done: every CTA's shared memory is there
|
||||
for (int i = t; i < CL_N * 64; i += CL_T) { // the histogram into slot `rank` of every CTA, as int4
|
||||
const int dst = i >> 6, c = i & 63;
|
||||
*cl_map(reinterpret_cast<int4*>(&hin[pass & 1][rank][0]) + c, (unsigned) dst) =
|
||||
reinterpret_cast<const int4*>(hl)[c];
|
||||
}
|
||||
for (int i = t; i < 256; i += CL_T) hloc[(pass + 1) & 1][i] = 0;
|
||||
cl_arrive();
|
||||
cl_wait();
|
||||
if (warp == 0) {
|
||||
// lane L owns digits 255 - 8L down to 248 - 8L; the digit holding the need-th cell from the top
|
||||
const int need = (int) width - above;
|
||||
int tot[8], part = 0;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 8; ++k) {
|
||||
const int d = 255 - 8 * lane - k;
|
||||
int v = 0;
|
||||
#pragma unroll
|
||||
for (int r = 0; r < CL_N; ++r) v += hin[pass & 1][r][d];
|
||||
tot[k] = v;
|
||||
part += v;
|
||||
}
|
||||
int incl = part;
|
||||
#pragma unroll
|
||||
for (int o = 1; o < 32; o <<= 1) {
|
||||
const int y = __shfl_up_sync(0xffffffffu, incl, o);
|
||||
if (lane >= o) incl += y;
|
||||
}
|
||||
const int excl = incl - part;
|
||||
const unsigned hit = __ballot_sync(0xffffffffu, excl < need && incl >= need);
|
||||
// no hit cannot happen (the cells total n_kv > width); the reference would then take digit 0
|
||||
const int who = hit ? __ffs(hit) - 1 : 31;
|
||||
if (lane == who) {
|
||||
int acc = excl, k = 0;
|
||||
for (; k < 7; ++k) {
|
||||
if (acc + tot[k] >= need) break;
|
||||
acc += tot[k];
|
||||
}
|
||||
s_prefix = prefix | ((uint32_t) (255 - 8 * lane - k) << shift);
|
||||
s_above = above + acc;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
prefix = s_prefix;
|
||||
above = s_above;
|
||||
}
|
||||
const uint32_t thr = prefix;
|
||||
const int64_t eq_budget = width - above; // cells at thr that fit, lowest index first
|
||||
// each thread a contiguous run of the CTA's keys; an odd run length keeps the shared reads conflict-free
|
||||
int seg = (n + CL_T - 1) / CL_T;
|
||||
if ((seg & 1) == 0 && seg > 0) ++seg;
|
||||
const int s0 = t * seg < n ? t * seg : n, s1 = s0 + seg < n ? s0 + seg : n;
|
||||
uint32_t gt = 0, eq = 0;
|
||||
for (int i = s0; i < s1; ++i) {
|
||||
const uint32_t k = keys[i], w = i == tail_i ? (uint32_t) wtail : (uint32_t) R;
|
||||
if (k > thr) gt += w;
|
||||
else if (k == thr) eq += w;
|
||||
}
|
||||
unsigned long long cta;
|
||||
const unsigned long long before =
|
||||
cl_excl_scan64((unsigned long long) gt | ((unsigned long long) eq << 32), s_warp, cta);
|
||||
if (t < CL_N) *cl_map(&cnt[rank], (unsigned) t) = cta;
|
||||
cl_arrive();
|
||||
cl_wait();
|
||||
int64_t gt_off = 0, eq_off = 0;
|
||||
for (unsigned r = 0; r < rank; ++r) {
|
||||
gt_off += (int64_t) (cnt[r] & 0xffffffffu);
|
||||
eq_off += (int64_t) (cnt[r] >> 32);
|
||||
}
|
||||
const int64_t eq_before = eq_off + (int64_t) (before >> 32);
|
||||
int64_t wpos = gt_off + (int64_t) (before & 0xffffffffu) + (eq_before < eq_budget ? eq_before : eq_budget);
|
||||
int64_t eq_left = eq_budget - eq_before;
|
||||
if (eq_left < 0) eq_left = 0;
|
||||
for (int i = s0; i < s1; ++i) {
|
||||
const uint32_t k = keys[i];
|
||||
const int w = i == tail_i ? wtail : R;
|
||||
const int64_t b = lo + i;
|
||||
if (k > thr) {
|
||||
for (int c = 0; c < w; ++c) out[wpos++] = (int32_t) (b * R + c);
|
||||
} else if (k == thr) {
|
||||
for (int c = 0; c < w && eq_left > 0; ++c, --eq_left) out[wpos++] = (int32_t) (b * R + c);
|
||||
}
|
||||
}
|
||||
#else
|
||||
(void) scores; (void) steps; (void) max_blocks; (void) cap; (void) ids;
|
||||
__trap();
|
||||
#endif
|
||||
}
|
||||
#endif // !__HIPCC__
|
||||
} // namespace
|
||||
|
||||
void qsa_block_scores(const float* pooled, const float* dead, const float* q_idx, const int32_t* steps, int64_t nq,
|
||||
@@ -745,12 +968,88 @@ void qsa_block_topk_ref(const float* scores, const int32_t* steps, int64_t nq, i
|
||||
if (e != cudaSuccess) { std::fprintf(stderr, "qsa_block_topk: %s\n", cudaGetErrorString(e)); std::exit(1); }
|
||||
}
|
||||
|
||||
bool qsa_block_topk_cluster(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) {
|
||||
#if defined(__HIPCC__)
|
||||
(void) scores; (void) steps; (void) nq; (void) max_blocks; (void) cap; (void) s; (void) ids; (void) stream;
|
||||
return false;
|
||||
#else
|
||||
if (nq <= 0) return true;
|
||||
if (s.idx_block != R || cap < qsa_selection_width(kTopkMaxCells, s) || nq > 65535 || max_blocks <= 0) return false;
|
||||
const size_t smem = (size_t) ((max_blocks + CL_N - 1) / CL_N) * sizeof(uint32_t);
|
||||
// Per device (a layer split runs on several): 1 the cluster kernel runs here, 2 it does not. It needs sm_90+ (the
|
||||
// card's, or STRATA_EMULATE_CC's) AND code built for it: a build with only older code JIT-compiles their PTX, whose
|
||||
// copy of this kernel is a trap - the function's PTX version says which. `opt`: the dynamic shared memory opted in.
|
||||
static int ok[64] = {};
|
||||
static size_t opt[64] = {};
|
||||
int dev = 0;
|
||||
if (cudaGetDevice(&dev) != cudaSuccess || dev < 0 || dev >= 64) { cudaGetLastError(); return false; }
|
||||
if (ok[dev] == 0) {
|
||||
int major = 0;
|
||||
cudaFuncAttributes fa{};
|
||||
const bool code = cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev) == cudaSuccess &&
|
||||
strata::cc_major_of(major) >= 9 &&
|
||||
cudaFuncGetAttributes(&fa, block_topk_cluster_kernel) == cudaSuccess && fa.ptxVersion >= 90 &&
|
||||
fa.binaryVersion >= 90;
|
||||
cudaGetLastError();
|
||||
ok[dev] = code ? 1 : 2;
|
||||
}
|
||||
if (ok[dev] != 1) return false;
|
||||
if (smem > opt[dev]) { // a larger capacity: opt in, and check that a cluster of CL_N such CTAs can be resident
|
||||
int clusters = 0;
|
||||
cudaLaunchConfig_t q{};
|
||||
cudaLaunchAttribute qa[1];
|
||||
qa[0].id = cudaLaunchAttributeClusterDimension;
|
||||
qa[0].val.clusterDim.x = CL_N;
|
||||
qa[0].val.clusterDim.y = 1;
|
||||
qa[0].val.clusterDim.z = 1;
|
||||
q.gridDim = dim3(CL_N, 1, 1);
|
||||
q.blockDim = dim3(CL_T, 1, 1);
|
||||
q.dynamicSmemBytes = smem;
|
||||
q.attrs = qa;
|
||||
q.numAttrs = 1;
|
||||
if (cudaFuncSetAttribute(block_topk_cluster_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int) smem) !=
|
||||
cudaSuccess ||
|
||||
cudaOccupancyMaxActiveClusters(&clusters, block_topk_cluster_kernel, &q) != cudaSuccess || clusters < 1) {
|
||||
cudaGetLastError();
|
||||
return false; // this capacity takes the one-CTA kernels; a smaller one may still fit
|
||||
}
|
||||
opt[dev] = smem;
|
||||
}
|
||||
cudaLaunchConfig_t cfg{};
|
||||
cudaLaunchAttribute at[1];
|
||||
at[0].id = cudaLaunchAttributeClusterDimension;
|
||||
at[0].val.clusterDim.x = CL_N;
|
||||
at[0].val.clusterDim.y = 1;
|
||||
at[0].val.clusterDim.z = 1;
|
||||
cfg.gridDim = dim3(CL_N, (unsigned) nq, 1);
|
||||
cfg.blockDim = dim3(CL_T, 1, 1);
|
||||
cfg.dynamicSmemBytes = smem;
|
||||
cfg.stream = (cudaStream_t) stream;
|
||||
cfg.attrs = at;
|
||||
cfg.numAttrs = 1;
|
||||
const cudaError_t e = cudaLaunchKernelEx(&cfg, block_topk_cluster_kernel, scores, steps, max_blocks, cap, ids);
|
||||
if (e != cudaSuccess) { std::fprintf(stderr, "qsa_block_topk cluster: %s\n", cudaGetErrorString(e)); std::exit(1); }
|
||||
return true;
|
||||
#endif
|
||||
}
|
||||
|
||||
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, int64_t active_blocks) {
|
||||
// keys in registers when every query's blocks fit (contexts up to ~135K cells); the same ids. STRATA_TOPK_OLD=1:
|
||||
// the kernel that reads them from memory on every pass
|
||||
static const bool old = std::getenv("STRATA_TOPK_OLD") != nullptr;
|
||||
if (nq <= 0) return;
|
||||
#if !defined(__HIPCC__)
|
||||
// sm_90+: a cluster of CL_N CTAs per query for the calls of a few queries (decode windows); the same ids.
|
||||
// STRATA_QSA_CLUSTER=0: the one-CTA kernels below
|
||||
static const bool cluster = [] {
|
||||
const char* v = std::getenv("STRATA_QSA_CLUSTER");
|
||||
return !v || std::atoi(v) != 0;
|
||||
}();
|
||||
if (cluster && !old && nq <= CL_MAXQ && qsa_block_topk_cluster(scores, steps, nq, max_blocks, cap, s, ids, stream))
|
||||
return;
|
||||
#endif
|
||||
// the blocks a query can have: the call's active count when the caller knows it (the prompt path), else the capacity.
|
||||
// Decode (no count) keeps the capacity rule and the original register width: nothing changes there.
|
||||
#if defined(__HIPCC__)
|
||||
|
||||
+151
-3
@@ -10,7 +10,8 @@
|
||||
// produces a valid token, so only a comparison at the distribution level can tell them apart; the parity test
|
||||
// does that against an independently computed distribution.
|
||||
//
|
||||
// `sampler_greedy_kernel` is the plain argmax, one block per token over the vocabulary. The sampled chain has
|
||||
// `sampler_greedy_kernel` is the plain argmax, one block per token over the vocabulary (on sm_90+ without penalties,
|
||||
// `sampler_greedy_cluster_kernel`: the same token from a cluster of 8 CTAs per row). The sampled chain has
|
||||
// three implementations that pick the same token, bit for bit:
|
||||
// - the SPLIT top_k (default): `sampler_split_part_kernel` cuts each row into 4,096-logit blocks over the whole
|
||||
// GPU, each keeps its own top_k, and `sampler_split_merge_kernel` merges those lists and runs the tail;
|
||||
@@ -20,6 +21,7 @@
|
||||
// The two new ones share `sampled_tail_warp` (top_p / min_p / temperature / draw on one warp).
|
||||
#include "strata/kernels/sampler.hpp"
|
||||
#include "strata/core/coupled_draft.hpp"
|
||||
#include "strata/core/emulate.hpp"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
@@ -166,6 +168,96 @@ __global__ void sampler_greedy_kernel(const float* __restrict__ logits, int n_vo
|
||||
}
|
||||
}
|
||||
|
||||
#if !defined(__HIPCC__)
|
||||
/// **THE SAME ARGMAX ON A THREAD-BLOCK CLUSTER (sm_90+, S19), FOR THE CALLS WITHOUT PENALTIES.** One block per
|
||||
/// token reads its 1 MB row (248,320 logits) on one SM with one load in flight per thread: latency-bound, 40 us per
|
||||
/// call on an RTX 5070, for every verify window and every draft step. Here a cluster of `kAmCtas` CTAs shares the
|
||||
/// row, each thread keeps four loads in flight, and the CTAs' results meet in CTA 0's shared memory (distributed
|
||||
/// shared memory): 6 us (decode_cluster_parity --bench).
|
||||
///
|
||||
/// The answer is a function of the row alone, so it is the one-block kernel's bit for bit: the LOWEST index whose value
|
||||
/// is the largest non-NaN value above -inf (each thread walks its elements in ascending order with a strict `>`, and
|
||||
/// every merge takes the larger value or, on equality, the smaller index - an order-free rule), and 0 when there is
|
||||
/// none (all -inf / NaN), as there. NaN never wins a `>`.
|
||||
///
|
||||
/// Barriers: a relaxed cluster arrive at entry, waited before the remote store (CTA 0 must be running); each CTA's
|
||||
/// result goes to slot `rank` of CTA 0, released by the next arrive; CTAs 1.. then exit (a cluster wait counts the
|
||||
/// threads that have not exited) and CTA 0 waits, reads its own slots, and writes the token.
|
||||
constexpr int kAmCtas = 8; // CTAs per token (the portable cluster size)
|
||||
constexpr int kAmThreads = 1024;
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
#define STRATA_AM_CLUSTER 1
|
||||
#else
|
||||
#define STRATA_AM_CLUSTER 0 // older targets: a trap, never launched (sample_greedy_cluster checks)
|
||||
#endif
|
||||
#if STRATA_AM_CLUSTER
|
||||
__device__ __forceinline__ void am_take(float s, int v, float& bv, int& best) {
|
||||
if (s > bv) { bv = s; best = v; }
|
||||
}
|
||||
__device__ __forceinline__ void am_merge(float ov, int oi, float& bv, int& best) {
|
||||
if (ov > bv || (ov == bv && oi < best)) { bv = ov; best = oi; }
|
||||
}
|
||||
#endif
|
||||
// grid (kAmCtas, n_tokens), cluster (kAmCtas, 1, 1), kAmThreads threads
|
||||
__global__ void __launch_bounds__(kAmThreads) sampler_greedy_cluster_kernel(const float* __restrict__ logits,
|
||||
int n_vocab, int* __restrict__ out) {
|
||||
#if STRATA_AM_CLUSTER
|
||||
__shared__ float sv[32];
|
||||
__shared__ int si[32];
|
||||
__shared__ float cv[kAmCtas]; // on CTA 0: each CTA's result
|
||||
__shared__ int ci[kAmCtas];
|
||||
unsigned rank;
|
||||
asm volatile("mov.u32 %0, %%cluster_ctarank;\n" : "=r"(rank));
|
||||
asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" ::: "memory");
|
||||
const float* l = logits + (size_t) blockIdx.y * n_vocab;
|
||||
float bv = __int_as_float(0xff800000); // -inf; `n_vocab` is "no candidate", as in sampler_greedy_kernel
|
||||
int best = n_vocab;
|
||||
constexpr int S = kAmCtas * kAmThreads;
|
||||
int v = (int) rank * kAmThreads + (int) threadIdx.x;
|
||||
for (; v + 3 * S < n_vocab; v += 4 * S) { // four independent loads, then taken in ascending order
|
||||
const float x0 = l[v], x1 = l[v + S], x2 = l[v + 2 * S], x3 = l[v + 3 * S];
|
||||
am_take(x0, v, bv, best);
|
||||
am_take(x1, v + S, bv, best);
|
||||
am_take(x2, v + 2 * S, bv, best);
|
||||
am_take(x3, v + 3 * S, bv, best);
|
||||
}
|
||||
for (; v < n_vocab; v += S) am_take(l[v], v, bv, best);
|
||||
for (int off = 16; off > 0; off >>= 1)
|
||||
am_merge(__shfl_down_sync(0xFFFFFFFFu, bv, off), __shfl_down_sync(0xFFFFFFFFu, best, off), bv, best);
|
||||
const int warp = (int) (threadIdx.x >> 5), lane = (int) (threadIdx.x & 31);
|
||||
if (lane == 0) { sv[warp] = bv; si[warp] = best; }
|
||||
__syncthreads();
|
||||
if (warp == 0) {
|
||||
bv = sv[lane];
|
||||
best = si[lane];
|
||||
for (int off = 16; off > 0; off >>= 1)
|
||||
am_merge(__shfl_down_sync(0xFFFFFFFFu, bv, off), __shfl_down_sync(0xFFFFFFFFu, best, off), bv, best);
|
||||
}
|
||||
asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); // CTA 0 runs
|
||||
if (threadIdx.x == 0) {
|
||||
uint64_t a;
|
||||
asm volatile("mapa.u64 %0, %1, %2;\n" : "=l"(a) : "l"((uint64_t) &cv[rank]), "r"(0u));
|
||||
*reinterpret_cast<float*>(a) = bv;
|
||||
asm volatile("mapa.u64 %0, %1, %2;\n" : "=l"(a) : "l"((uint64_t) &ci[rank]), "r"(0u));
|
||||
*reinterpret_cast<int*>(a) = best;
|
||||
}
|
||||
asm volatile("barrier.cluster.arrive.release.aligned;\n" ::: "memory");
|
||||
if (rank != 0) return;
|
||||
asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory");
|
||||
if (warp == 0) {
|
||||
bv = lane < kAmCtas ? cv[lane] : __int_as_float(0xff800000);
|
||||
best = lane < kAmCtas ? ci[lane] : n_vocab;
|
||||
for (int off = 16; off > 0; off >>= 1)
|
||||
am_merge(__shfl_down_sync(0xFFFFFFFFu, bv, off), __shfl_down_sync(0xFFFFFFFFu, best, off), bv, best);
|
||||
if (lane == 0) out[blockIdx.y] = (best < n_vocab) ? best : 0;
|
||||
}
|
||||
#else
|
||||
(void) logits; (void) n_vocab; (void) out;
|
||||
__trap();
|
||||
#endif
|
||||
}
|
||||
#endif // !__HIPCC__
|
||||
|
||||
/// **THE SAMPLED PATH, ONE BLOCK PER TOKEN.** The kernel below replaced a version that ran the whole chain
|
||||
/// in ONE THREAD per token (`<<<ceil(T/64), 64>>>`, so a 4-token window fielded four threads): `top_k` alone
|
||||
/// was `k` sequential scans of the vocabulary with an inner sweep over the already-taken list - 20 x 248,320
|
||||
@@ -833,6 +925,55 @@ int2* split_scratch(void* stream, size_t entries) {
|
||||
|
||||
} // namespace
|
||||
|
||||
bool sample_greedy_cluster(const float* logits, int n_tokens, int n_vocab, int* out, void* stream) {
|
||||
#if defined(__HIPCC__)
|
||||
(void) logits; (void) n_tokens; (void) n_vocab; (void) out; (void) stream;
|
||||
return false;
|
||||
#else
|
||||
if (n_tokens <= 0 || n_vocab <= 0) return true;
|
||||
if (n_tokens > 65535) return false;
|
||||
// Per device (a layer split runs on several): 1 the cluster kernel runs here, 2 it does not - sm_90+ (the card's,
|
||||
// or STRATA_EMULATE_CC's), code built for it (an older build's PTX holds a trap: the PTX version says which), and
|
||||
// room for one cluster of kAmCtas CTAs.
|
||||
static int ok[64] = {};
|
||||
int dev = 0;
|
||||
if (cudaGetDevice(&dev) != cudaSuccess || dev < 0 || dev >= 64) { cudaGetLastError(); return false; }
|
||||
cudaLaunchConfig_t cfg{};
|
||||
cudaLaunchAttribute at[1];
|
||||
at[0].id = cudaLaunchAttributeClusterDimension;
|
||||
at[0].val.clusterDim.x = kAmCtas;
|
||||
at[0].val.clusterDim.y = 1;
|
||||
at[0].val.clusterDim.z = 1;
|
||||
cfg.gridDim = dim3(kAmCtas, 1, 1);
|
||||
cfg.blockDim = dim3(kAmThreads, 1, 1);
|
||||
cfg.dynamicSmemBytes = 0;
|
||||
cfg.stream = (cudaStream_t) stream;
|
||||
cfg.attrs = at;
|
||||
cfg.numAttrs = 1;
|
||||
if (ok[dev] == 0) {
|
||||
int major = 0, clusters = 0;
|
||||
cudaFuncAttributes fa{};
|
||||
const bool runs = cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev) == cudaSuccess &&
|
||||
strata::cc_major_of(major) >= 9 &&
|
||||
cudaFuncGetAttributes(&fa, sampler_greedy_cluster_kernel) == cudaSuccess &&
|
||||
fa.ptxVersion >= 90 && fa.binaryVersion >= 90 &&
|
||||
cudaOccupancyMaxActiveClusters(&clusters, sampler_greedy_cluster_kernel, &cfg) ==
|
||||
cudaSuccess &&
|
||||
clusters >= 1;
|
||||
cudaGetLastError();
|
||||
ok[dev] = runs ? 1 : 2;
|
||||
}
|
||||
if (ok[dev] != 1) return false;
|
||||
cfg.gridDim = dim3(kAmCtas, (unsigned) n_tokens, 1);
|
||||
const cudaError_t e = cudaLaunchKernelEx(&cfg, sampler_greedy_cluster_kernel, logits, n_vocab, out);
|
||||
if (e != cudaSuccess) {
|
||||
std::fprintf(stderr, "sample_tokens cluster launch: %s\n", cudaGetErrorString(e));
|
||||
std::exit(1);
|
||||
}
|
||||
return true;
|
||||
#endif
|
||||
}
|
||||
|
||||
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) {
|
||||
if (n_tokens <= 0 || n_vocab <= 0) return;
|
||||
@@ -845,10 +986,17 @@ void sample_tokens(const float* logits, int n_tokens, int n_vocab, const int* hi
|
||||
? (unsigned) ((n_vocab + 31) / 32) * sizeof(unsigned) // the penalty bitmap
|
||||
: 0;
|
||||
if (p.greedy || p.temperature <= 0.0f) {
|
||||
// Without penalties (shmem == 0: no window) on sm_90+, a cluster of CTAs per token - the same token; see
|
||||
// `sampler_greedy_cluster_kernel`. STRATA_ARGMAX_MULTI=0: always the one-block kernel.
|
||||
static const bool multi = [] {
|
||||
const char* v = std::getenv("STRATA_ARGMAX_MULTI");
|
||||
return !v || std::atoi(v) != 0;
|
||||
}();
|
||||
// One block per token, 1,024 threads over the vocabulary. See `sampler_greedy_kernel`.
|
||||
const int gthreads = 1024;
|
||||
sampler_greedy_kernel<<<(unsigned) n_tokens, gthreads, shmem, (cudaStream_t) stream>>>(
|
||||
logits, n_vocab, history, history_len, p, p.penalty_last_n, p.penalty_last_n, out);
|
||||
if (!(multi && shmem == 0 && sample_greedy_cluster(logits, n_tokens, n_vocab, out, stream)))
|
||||
sampler_greedy_kernel<<<(unsigned) n_tokens, gthreads, shmem, (cudaStream_t) stream>>>(
|
||||
logits, n_vocab, history, history_len, p, p.penalty_last_n, p.penalty_last_n, out);
|
||||
} else if (sampled_path() == SampledPath::Old) {
|
||||
// The same block-per-token shape: the selection's k argmax rounds reduce inside the block. See
|
||||
// `sampler_kernel`'s header for what the old one-thread-per-token launch cost.
|
||||
|
||||
@@ -0,0 +1,421 @@
|
||||
// src/kernels/decode_cluster_parity.cpp - S19: the decode kernels on thread-block clusters (sm_90+) against the ones
|
||||
// they replace, bit for bit (GPU, synthetic, no model).
|
||||
//
|
||||
// QSA top-k : qsa_block_topk_cluster vs qsa_block_topk's one-CTA dispatch (the register kernel up to 33,792 blocks,
|
||||
// block_topk_kernel above) and vs qsa_block_topk_ref - every id of every row of the `cap`-wide buffer,
|
||||
// including the entries past the width that neither may write.
|
||||
// argmax : sample_greedy_cluster vs sample_tokens' one-block greedy kernel.
|
||||
//
|
||||
// Inputs: random scores like a real indexer's, and adversarial ones - all equal, a few distinct values (ties at the
|
||||
// threshold across the CTAs of a cluster), NaN, +-inf, -0.0 against +0.0, denormals, monotone rows, all NaN; contexts
|
||||
// from below the selection width (the identity) to 262,144 cells, every tail weight (n_kv mod 4), 1-16 queries; and
|
||||
// one CUDA graph that captures both cluster launches, replayed with new inputs.
|
||||
//
|
||||
// decode_cluster_parity --selftest parity (exit 0 = identical everywhere); skipped (exit 0) below sm_90
|
||||
// decode_cluster_parity --bench per-call times, old vs new, at 32K / 128K / 262K
|
||||
//
|
||||
// The old paths are reached through the public dispatchers with STRATA_QSA_CLUSTER=0 / STRATA_ARGMAX_MULTI=0, set
|
||||
// here before their first call (they read the variables once).
|
||||
#include "strata/kernels/qsa.hpp"
|
||||
#include "strata/kernels/qsa_select.hpp"
|
||||
#include "strata/kernels/sampler.hpp"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <limits>
|
||||
#include <random>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace k = strata::kernels;
|
||||
|
||||
namespace {
|
||||
|
||||
void ck(cudaError_t e, const char* what) {
|
||||
if (e != cudaSuccess) {
|
||||
std::fprintf(stderr, "%s: %s\n", what, cudaGetErrorString(e));
|
||||
std::exit(2);
|
||||
}
|
||||
}
|
||||
|
||||
void set_env(const char* name, const char* value) {
|
||||
#if defined(_WIN32)
|
||||
_putenv_s(name, value);
|
||||
#else
|
||||
setenv(name, value, 1);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename T> struct Dev {
|
||||
T* p = nullptr;
|
||||
size_t n = 0;
|
||||
explicit Dev(size_t count) : n(count) { ck(cudaMalloc(&p, count * sizeof(T) + 64), "cudaMalloc"); }
|
||||
~Dev() { cudaFree(p); }
|
||||
void put(const std::vector<T>& h) {
|
||||
ck(cudaMemcpy(p, h.data(), h.size() * sizeof(T), cudaMemcpyHostToDevice), "up");
|
||||
}
|
||||
std::vector<T> get() const {
|
||||
std::vector<T> h(n);
|
||||
ck(cudaMemcpy(h.data(), p, n * sizeof(T), cudaMemcpyDeviceToHost), "down");
|
||||
return h;
|
||||
}
|
||||
};
|
||||
|
||||
const float kNaN = std::numeric_limits<float>::quiet_NaN(), kInf = std::numeric_limits<float>::infinity();
|
||||
|
||||
// ---- QSA top-k fixtures
|
||||
enum Pattern { kRandom, kAllEqual, kFewValues, kNaNs, kInfs, kSignedZero, kDenormal, kAscending, kDescending, kAllNaN,
|
||||
kQuantized, kPatterns };
|
||||
const char* pattern_name(int p) {
|
||||
static const char* n[] = {"random", "all-equal", "few-values", "nan", "inf", "signed-zero", "denormal", "ascending",
|
||||
"descending", "all-nan", "quantized"};
|
||||
return n[p];
|
||||
}
|
||||
|
||||
// one query's row of block scores (blocks 0..n_bid); the tail block gets the +1e9 the score kernel gives it
|
||||
void fill_row(float* row, int64_t n_kv, int64_t n_bid, int pattern, std::mt19937& rng) {
|
||||
std::normal_distribution<float> nd(0.f, 1.f);
|
||||
std::uniform_int_distribution<int> pick(0, 99);
|
||||
for (int64_t b = 0; b <= n_bid; ++b) {
|
||||
float v = 0.f;
|
||||
const float g = nd(rng);
|
||||
switch (pattern) {
|
||||
case kRandom: v = std::fmax(0.f, g) * 3.f + std::fmax(0.f, nd(rng)); break; // relu sums: exact zeros too
|
||||
case kAllEqual: v = 1.f; break;
|
||||
case kFewValues: v = (float) (pick(rng) % 4) * 0.5f; break;
|
||||
case kNaNs: v = pick(rng) < 15 ? kNaN : std::fabs(g); break;
|
||||
case kInfs: {
|
||||
const int r = pick(rng);
|
||||
v = r < 5 ? kInf : r < 15 ? -kInf : r < 30 ? -std::fabs(g) : g;
|
||||
break;
|
||||
}
|
||||
case kSignedZero: v = pick(rng) < 50 ? -0.0f : (pick(rng) < 50 ? 0.0f : -std::fabs(g) * 1e-3f); break;
|
||||
case kDenormal: v = (float) (pick(rng) % 7) * 1e-41f; break;
|
||||
case kAscending: v = (float) b; break;
|
||||
case kDescending: v = (float) (n_bid - b); break;
|
||||
case kAllNaN: v = kNaN; break;
|
||||
case kQuantized: v = std::round(std::fabs(g) * 8.f) / 8.f; break;
|
||||
}
|
||||
row[b] = v;
|
||||
}
|
||||
if (n_kv % 4 != 0 && pattern != kAllNaN && pattern != kAllEqual) row[n_bid] += 1e9f;
|
||||
}
|
||||
|
||||
struct TopkCase {
|
||||
int64_t ctx; // the last query's n_kv
|
||||
int64_t capacity; // --max-context: max_blocks = capacity / 4 + 2 (the engine's rule)
|
||||
int nq;
|
||||
int pattern;
|
||||
};
|
||||
|
||||
int run_topk_case(const TopkCase& c, std::mt19937& rng, bool verbose) {
|
||||
const k::QsaShapes s = k::qsa_real_shapes();
|
||||
const int64_t max_blocks = c.capacity / 4 + 2, cap = k::qsa_selection_width(k::kTopkMaxCells, s);
|
||||
std::vector<float> sc((size_t) (c.nq * max_blocks), 0.f);
|
||||
std::vector<int32_t> steps((size_t) c.nq * k::kStepCount);
|
||||
for (int i = 0; i < c.nq; ++i) {
|
||||
const int64_t n_kv = c.ctx - (c.nq - 1 - i); // the window's queries, one cell apart, the last at ctx
|
||||
int32_t* st = steps.data() + (size_t) i * k::kStepCount;
|
||||
st[k::kStepPos] = (int32_t) (n_kv - 1);
|
||||
st[k::kStepNKv] = (int32_t) n_kv;
|
||||
st[k::kStepNBid] = (int32_t) (n_kv / s.idx_block);
|
||||
st[k::kStepWidth] = (int32_t) k::qsa_selection_width(n_kv, s);
|
||||
fill_row(sc.data() + (size_t) i * max_blocks, n_kv, n_kv / s.idx_block, c.pattern, rng);
|
||||
}
|
||||
Dev<float> d_sc(sc.size());
|
||||
const size_t n_ids = (size_t) (c.nq * cap);
|
||||
Dev<int32_t> d_st(steps.size()), d_new(n_ids), d_old(n_ids), d_ref(n_ids);
|
||||
d_sc.put(sc);
|
||||
d_st.put(steps);
|
||||
ck(cudaMemset(d_new.p, 0xA5, d_new.n * 4), "memset");
|
||||
ck(cudaMemset(d_old.p, 0xA5, d_old.n * 4), "memset");
|
||||
ck(cudaMemset(d_ref.p, 0xA5, d_ref.n * 4), "memset");
|
||||
if (!k::qsa_block_topk_cluster(d_sc.p, d_st.p, c.nq, max_blocks, cap, s, d_new.p, nullptr)) {
|
||||
std::printf("FAIL topk ctx %lld cap %lld nq %d: the cluster kernel did not run\n", (long long) c.ctx,
|
||||
(long long) c.capacity, c.nq);
|
||||
return 1;
|
||||
}
|
||||
k::qsa_block_topk(d_sc.p, d_st.p, c.nq, max_blocks, cap, s, d_old.p, nullptr);
|
||||
k::qsa_block_topk_ref(d_sc.p, d_st.p, c.nq, max_blocks, cap, s, d_ref.p, nullptr);
|
||||
ck(cudaDeviceSynchronize(), "topk sync");
|
||||
const std::vector<int32_t> a = d_new.get(), b = d_old.get(), r = d_ref.get();
|
||||
int64_t diff_old = 0, diff_ref = 0, first = -1;
|
||||
for (size_t i = 0; i < a.size(); ++i) {
|
||||
if (a[i] != b[i]) { ++diff_old; if (first < 0) first = (int64_t) i; }
|
||||
if (a[i] != r[i]) ++diff_ref;
|
||||
}
|
||||
const bool ok = diff_old == 0 && diff_ref == 0;
|
||||
if (!ok || verbose)
|
||||
std::printf("%s topk ctx %7lld cap %7lld nq %2d %-11s: %lld ids differ from the one-CTA dispatch, %lld from "
|
||||
"the reference%s\n", ok ? "ok " : "FAIL", (long long) c.ctx, (long long) c.capacity, c.nq,
|
||||
pattern_name(c.pattern), (long long) diff_old, (long long) diff_ref,
|
||||
first >= 0 ? (" (first at " + std::to_string(first) + ")").c_str() : "");
|
||||
return ok ? 0 : 1;
|
||||
}
|
||||
|
||||
// ---- argmax fixtures
|
||||
enum AmPattern { aRandom, aTies, aAllEqual, aAllNegInf, aAllNaN, aNaNMax, aPosInf, aSignedZero, aLast, aFirst,
|
||||
aDenormal, aTiesAcross, aPatterns };
|
||||
const char* am_name(int p) {
|
||||
static const char* n[] = {"random", "ties", "all-equal", "all--inf", "all-nan", "nan+max", "+inf", "signed-zero",
|
||||
"max-last", "max-first", "denormal", "ties-across"};
|
||||
return n[p];
|
||||
}
|
||||
|
||||
void fill_logits(float* l, int nv, int pattern, std::mt19937& rng) {
|
||||
std::normal_distribution<float> nd(0.f, 3.f);
|
||||
std::uniform_int_distribution<int> at(0, nv - 1);
|
||||
for (int v = 0; v < nv; ++v) l[v] = nd(rng);
|
||||
switch (pattern) {
|
||||
case aRandom: break;
|
||||
case aTies: { const int n = 1 + at(rng) % 9; for (int i = 0; i < n; ++i) l[at(rng)] = 100.f; break; }
|
||||
case aAllEqual: for (int v = 0; v < nv; ++v) l[v] = 0.25f; break;
|
||||
case aAllNegInf: for (int v = 0; v < nv; ++v) l[v] = -kInf; break;
|
||||
case aAllNaN: for (int v = 0; v < nv; ++v) l[v] = kNaN; break;
|
||||
case aNaNMax:
|
||||
for (int v = 0; v < nv; ++v) if (at(rng) % 3 == 0) l[v] = kNaN;
|
||||
l[at(rng)] = 50.f;
|
||||
l[at(rng)] = 50.f;
|
||||
break;
|
||||
case aPosInf: for (int i = 0; i < 4; ++i) l[at(rng)] = kInf; break;
|
||||
case aSignedZero:
|
||||
for (int v = 0; v < nv; ++v) l[v] = -std::fabs(l[v]) - 1.f;
|
||||
for (int i = 0; i < 6; ++i) l[at(rng)] = (i & 1) ? 0.0f : -0.0f;
|
||||
break;
|
||||
case aLast: l[nv - 1] = 1e30f; break;
|
||||
case aFirst: l[0] = 1e30f; if (nv > 1) l[nv - 1] = 1e30f; break;
|
||||
case aDenormal: for (int v = 0; v < nv; ++v) l[v] = (float) (at(rng) % 5) * 1e-42f - 1e-30f * (v & 1); break;
|
||||
case aTiesAcross: // the same maximum in several CTAs' and threads' ranges (8 x 1024 stride), plus -inf, NaN
|
||||
for (int v = 0; v < nv; v += 1 + at(rng) % 9000) l[v] = 7.f;
|
||||
for (int v = 3; v < nv; v += 11) l[v] = kNaN;
|
||||
for (int v = 5; v < nv; v += 13) l[v] = -kInf;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
int run_argmax_case(int nv, int T, int pattern, std::mt19937& rng, bool verbose) {
|
||||
std::vector<float> l((size_t) nv * T);
|
||||
for (int t = 0; t < T; ++t) fill_logits(l.data() + (size_t) t * nv, nv, t % 2 == 0 ? pattern : aRandom, rng);
|
||||
Dev<float> d_l(l.size());
|
||||
Dev<int> d_new((size_t) T), d_old((size_t) T);
|
||||
d_l.put(l);
|
||||
ck(cudaMemset(d_new.p, 0xA5, (size_t) T * 4), "memset");
|
||||
ck(cudaMemset(d_old.p, 0x5A, (size_t) T * 4), "memset");
|
||||
k::SamplerParams sp;
|
||||
sp.greedy = true;
|
||||
sp.temperature = 0.f;
|
||||
if (!k::sample_greedy_cluster(d_l.p, T, nv, d_new.p, nullptr)) {
|
||||
std::printf("FAIL argmax nv %d T %d: the cluster kernel did not run\n", nv, T);
|
||||
return 1;
|
||||
}
|
||||
k::sample_tokens(d_l.p, T, nv, nullptr, 0, sp, d_old.p, nullptr);
|
||||
ck(cudaDeviceSynchronize(), "argmax sync");
|
||||
const std::vector<int> a = d_new.get(), b = d_old.get();
|
||||
// and an independent host answer: the lowest index of the largest non-NaN value above -inf, else 0
|
||||
int bad = 0;
|
||||
for (int t = 0; t < T; ++t) {
|
||||
const float* r = l.data() + (size_t) t * nv;
|
||||
int hb = -1;
|
||||
for (int v = 0; v < nv; ++v)
|
||||
if (r[v] > -kInf && (hb < 0 || r[v] > r[hb])) hb = v;
|
||||
if (hb < 0) hb = 0;
|
||||
if (a[(size_t) t] != b[(size_t) t] || a[(size_t) t] != hb) ++bad;
|
||||
}
|
||||
if (bad || verbose)
|
||||
std::printf("%s argmax nv %6d T %d %-11s: %d rows differ (new %d, old %d)\n", bad ? "FAIL" : "ok ", nv, T,
|
||||
am_name(pattern), bad, a[0], b[0]);
|
||||
return bad ? 1 : 0;
|
||||
}
|
||||
|
||||
// both cluster launches captured in one graph, replayed over new inputs written into the same buffers
|
||||
int run_graph_case(std::mt19937& rng) {
|
||||
const k::QsaShapes s = k::qsa_real_shapes();
|
||||
const int nq = 5, nv = 248320;
|
||||
const int64_t capacity = 131072 + 272, max_blocks = capacity / 4 + 2;
|
||||
const int64_t cap = k::qsa_selection_width(k::kTopkMaxCells, s);
|
||||
Dev<float> d_sc((size_t) (nq * max_blocks)), d_l((size_t) nv * nq);
|
||||
Dev<int32_t> d_st((size_t) nq * k::kStepCount), d_new((size_t) (nq * cap)), d_old((size_t) (nq * cap));
|
||||
Dev<int> a_new((size_t) nq), a_old((size_t) nq);
|
||||
ck(cudaMemset(d_new.p, 0, d_new.n * 4), "memset"); // the entries past a row's width: written by neither
|
||||
ck(cudaMemset(d_old.p, 0, d_old.n * 4), "memset");
|
||||
cudaStream_t cs;
|
||||
ck(cudaStreamCreateWithFlags(&cs, cudaStreamNonBlocking), "stream");
|
||||
cudaGraph_t g = nullptr;
|
||||
cudaGraphExec_t ge = nullptr;
|
||||
ck(cudaStreamBeginCapture(cs, cudaStreamCaptureModeThreadLocal), "begin capture");
|
||||
const bool l1 = k::qsa_block_topk_cluster(d_sc.p, d_st.p, nq, max_blocks, cap, s, d_new.p, cs);
|
||||
const bool l2 = k::sample_greedy_cluster(d_l.p, nq, nv, a_new.p, cs);
|
||||
ck(cudaStreamEndCapture(cs, &g), "end capture");
|
||||
if (!l1 && !l2) { // no cluster code here (a card or a build below sm_90): nothing to compare
|
||||
if (g) cudaGraphDestroy(g);
|
||||
cudaStreamDestroy(cs);
|
||||
return -1;
|
||||
}
|
||||
if (!l1 || !l2) { std::printf("FAIL graph: one cluster kernel did not run\n"); return 1; }
|
||||
ck(cudaGraphInstantiate(&ge, g, 0), "instantiate");
|
||||
int fails = 0;
|
||||
for (int rep = 0; rep < 6; ++rep) {
|
||||
const int64_t ctx = 100000 + rep * 4001 + rep % 4;
|
||||
std::vector<float> sc((size_t) (nq * max_blocks), 0.f), l((size_t) nv * nq);
|
||||
std::vector<int32_t> steps((size_t) nq * k::kStepCount);
|
||||
for (int i = 0; i < nq; ++i) {
|
||||
const int64_t n_kv = ctx - (nq - 1 - i);
|
||||
int32_t* st = steps.data() + (size_t) i * k::kStepCount;
|
||||
st[k::kStepPos] = (int32_t) (n_kv - 1);
|
||||
st[k::kStepNKv] = (int32_t) n_kv;
|
||||
st[k::kStepNBid] = (int32_t) (n_kv / 4);
|
||||
st[k::kStepWidth] = (int32_t) k::qsa_selection_width(n_kv, s);
|
||||
fill_row(sc.data() + (size_t) i * max_blocks, n_kv, n_kv / 4, rep % 2 ? kFewValues : kRandom, rng);
|
||||
fill_logits(l.data() + (size_t) i * nv, nv, rep % 2 ? aTiesAcross : aRandom, rng);
|
||||
}
|
||||
d_sc.put(sc);
|
||||
d_st.put(steps);
|
||||
d_l.put(l);
|
||||
ck(cudaGraphLaunch(ge, cs), "graph launch");
|
||||
k::qsa_block_topk(d_sc.p, d_st.p, nq, max_blocks, cap, s, d_old.p, cs);
|
||||
k::SamplerParams sp;
|
||||
sp.greedy = true;
|
||||
k::sample_tokens(d_l.p, nq, nv, nullptr, 0, sp, a_old.p, cs);
|
||||
ck(cudaStreamSynchronize(cs), "graph sync");
|
||||
const auto x = d_new.get(), y = d_old.get();
|
||||
const auto u = a_new.get(), v = a_old.get();
|
||||
int64_t d1 = 0, d2 = 0;
|
||||
for (size_t i = 0; i < x.size(); ++i) d1 += x[i] != y[i];
|
||||
for (size_t i = 0; i < u.size(); ++i) d2 += u[i] != v[i];
|
||||
if (d1 || d2) ++fails;
|
||||
std::printf("%s graph replay %d (ctx %lld): %lld ids, %lld tokens differ\n", d1 || d2 ? "FAIL" : "ok ", rep,
|
||||
(long long) ctx, (long long) d1, (long long) d2);
|
||||
}
|
||||
cudaGraphExecDestroy(ge);
|
||||
cudaGraphDestroy(g);
|
||||
cudaStreamDestroy(cs);
|
||||
return fails;
|
||||
}
|
||||
|
||||
// back-to-back launches on `cs` (the engine's are too: one stream, inside a graph)
|
||||
template <typename F> float time_us(cudaStream_t cs, F&& f, int reps = 300) {
|
||||
cudaEvent_t a, b;
|
||||
cudaEventCreate(&a);
|
||||
cudaEventCreate(&b);
|
||||
for (int i = 0; i < 5; ++i) f();
|
||||
cudaEventRecord(a, cs);
|
||||
for (int i = 0; i < reps; ++i) f();
|
||||
cudaEventRecord(b, cs);
|
||||
ck(cudaEventSynchronize(b), "time");
|
||||
float ms = 0.f;
|
||||
cudaEventElapsedTime(&ms, a, b);
|
||||
cudaEventDestroy(a);
|
||||
cudaEventDestroy(b);
|
||||
return ms * 1000.f / (float) reps;
|
||||
}
|
||||
|
||||
int bench() {
|
||||
std::mt19937 rng(11);
|
||||
const k::QsaShapes s = k::qsa_real_shapes();
|
||||
const int64_t cap = k::qsa_selection_width(k::kTopkMaxCells, s);
|
||||
cudaStream_t cs;
|
||||
ck(cudaStreamCreateWithFlags(&cs, cudaStreamNonBlocking), "stream");
|
||||
std::printf("QSA top-k, per call (us); capacity = context + 272 as the speed runs set --max-context\n");
|
||||
for (const int64_t ctx : {32768LL, 131072LL, 262144LL - 272}) {
|
||||
for (const int nq : {1, 5}) {
|
||||
const int64_t capacity = ctx + 272, max_blocks = capacity / 4 + 2;
|
||||
std::vector<float> sc((size_t) (nq * max_blocks), 0.f);
|
||||
std::vector<int32_t> steps((size_t) nq * k::kStepCount);
|
||||
for (int i = 0; i < nq; ++i) {
|
||||
const int64_t n_kv = ctx - (nq - 1 - i);
|
||||
int32_t* st = steps.data() + (size_t) i * k::kStepCount;
|
||||
st[k::kStepPos] = (int32_t) (n_kv - 1);
|
||||
st[k::kStepNKv] = (int32_t) n_kv;
|
||||
st[k::kStepNBid] = (int32_t) (n_kv / 4);
|
||||
st[k::kStepWidth] = (int32_t) k::qsa_selection_width(n_kv, s);
|
||||
fill_row(sc.data() + (size_t) i * max_blocks, n_kv, n_kv / 4, kRandom, rng);
|
||||
}
|
||||
Dev<float> d_sc(sc.size());
|
||||
Dev<int32_t> d_st(steps.size()), d_ids((size_t) (nq * cap));
|
||||
d_sc.put(sc);
|
||||
d_st.put(steps);
|
||||
const float* sp = d_sc.p;
|
||||
const int32_t* st = d_st.p;
|
||||
int32_t* ids = d_ids.p;
|
||||
const float t_old = time_us(cs, [&] { k::qsa_block_topk(sp, st, nq, max_blocks, cap, s, ids, cs); });
|
||||
const float t_ref = time_us(cs, [&] { k::qsa_block_topk_ref(sp, st, nq, max_blocks, cap, s, ids, cs); });
|
||||
const float t_new = time_us(cs, [&] { k::qsa_block_topk_cluster(sp, st, nq, max_blocks, cap, s, ids, cs); });
|
||||
std::printf(" ctx %6lld nq %d: one-CTA dispatch %7.1f (ref %7.1f) cluster %6.1f -> %.1fx\n",
|
||||
(long long) ctx, nq, t_old, t_ref, t_new, t_old / t_new);
|
||||
}
|
||||
}
|
||||
std::printf("argmax, per call (us)\n");
|
||||
for (const int nv : {248320, 32768}) {
|
||||
for (const int T : {1, 5}) {
|
||||
std::vector<float> l((size_t) nv * T);
|
||||
for (int t = 0; t < T; ++t) fill_logits(l.data() + (size_t) t * nv, nv, aRandom, rng);
|
||||
Dev<float> d_l(l.size());
|
||||
Dev<int> d_o((size_t) T);
|
||||
d_l.put(l);
|
||||
k::SamplerParams sp;
|
||||
sp.greedy = true;
|
||||
const float t_old = time_us(cs, [&] { k::sample_tokens(d_l.p, T, nv, nullptr, 0, sp, d_o.p, cs); });
|
||||
const float t_new = time_us(cs, [&] { k::sample_greedy_cluster(d_l.p, T, nv, d_o.p, cs); });
|
||||
std::printf(" n_vocab %6d T %d: one block %6.1f cluster %5.1f -> %.1fx\n", nv, T, t_old, t_new,
|
||||
t_old / t_new);
|
||||
}
|
||||
}
|
||||
cudaStreamDestroy(cs);
|
||||
return 0;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
set_env("STRATA_QSA_CLUSTER", "0"); // the dispatchers: the one-CTA kernels (the cluster ones called directly)
|
||||
set_env("STRATA_ARGMAX_MULTI", "0");
|
||||
const bool do_bench = argc > 1 && std::strcmp(argv[1], "--bench") == 0;
|
||||
const bool verbose = argc > 2 && std::strcmp(argv[2], "-v") == 0;
|
||||
int dev = 0, major = 0;
|
||||
ck(cudaGetDevice(&dev), "device");
|
||||
ck(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev), "cc");
|
||||
std::mt19937 rng(20261002);
|
||||
// first, so that both kernels' first calls (their per-device checks and opt-ins) happen inside a stream capture,
|
||||
// as in the engine's decode graphs
|
||||
const int gfails = run_graph_case(rng);
|
||||
if (gfails < 0) {
|
||||
std::printf("SKIP: no thread-block clusters on this card / build (compute capability %d)\n", major);
|
||||
return 0;
|
||||
}
|
||||
if (do_bench) return bench();
|
||||
int fails = 0, cases = 0;
|
||||
// contexts: below the width (identity), at it, just past it, odd tails, the register kernel's range, past it
|
||||
// (block_topk_kernel), up to 262,144; capacities: tight (the speed runs) and 262,144 (a long --max-context)
|
||||
const int64_t ctxs[] = {7, 2050, 2051, 2052, 2053, 2054, 2055, 3001, 9000, 32768, 32769, 65538, 100003,
|
||||
131072, 135170, 135171, 200001, 262142, 262143, 262144};
|
||||
for (const int64_t ctx : ctxs) {
|
||||
for (const int64_t capacity : {ctx + 272, (int64_t) 262144 + 272}) {
|
||||
if (capacity < ctx) continue;
|
||||
for (int p = 0; p < kPatterns; ++p) {
|
||||
const int nq = (p % 3 == 0) ? 1 : (p % 3 == 1 ? 5 : 16);
|
||||
if (nq > ctx) continue;
|
||||
fails += run_topk_case({ctx, capacity, nq, p}, rng, verbose);
|
||||
++cases;
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("QSA top-k: %d cases, %d failed\n", cases, fails);
|
||||
int afails = 0, acases = 0;
|
||||
for (const int nv : {248320, 248321, 248319, 32771, 24576, 8193, 1000, 33, 1}) {
|
||||
for (int p = 0; p < aPatterns; ++p) {
|
||||
for (const int T : {1, 5}) {
|
||||
afails += run_argmax_case(nv, T, p, rng, verbose);
|
||||
++acases;
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("argmax: %d cases, %d failed\n", acases, afails);
|
||||
const int total = fails + afails + gfails;
|
||||
std::printf("%s\n", total == 0 ? "PASS: the cluster kernels are bitwise identical" : "FAIL");
|
||||
return total == 0 ? 0 : 1;
|
||||
}
|
||||
@@ -0,0 +1,459 @@
|
||||
// src/prefill/moe_fused.cu - see include/strata/prefill/moe_fused.hpp (#136: the Q2_0 pack's prompt experts on
|
||||
// Strata's own int8 tensor-core kernels, with the grouping on the GPU).
|
||||
//
|
||||
// The arithmetic. A Q2_0 block is 64 weights w = d_w * (q - 1), q in 0..3, four codes per byte, weight j of a block
|
||||
// in byte j / 4 at bits 2 * (j % 4). An activation block of 32 values is x = d_x * a with a = round(x / d_x) in
|
||||
// -127..127. Over one 32-value half-block
|
||||
// sum w x = d_w * d_x * (sum q a - sum a),
|
||||
// so the int8 product takes the codes as they are (0..3 is a valid s8) and the offset is the activations' own sum,
|
||||
// stored with the block. The dot product (|sum q a| <= 32 * 3 * 127 < 2^22) becomes a float exactly with one
|
||||
// integer add: as_float(0x4B400000 + dot) = 1.5 * 2^23 + dot, and the block stores c = -(1.5 * 2^23 + sum a), so
|
||||
// as_float(0x4B400000 + dot) + c = dot - sum a with no rounding. Per output and 64-block that is
|
||||
// acc += d_w * (d_x0 * t0 + d_x1 * t1) (t = dot - sum a of each half)
|
||||
// - two integer adds, two float adds and three multiply-adds.
|
||||
//
|
||||
// The K order. mma.sync m16n8k32 takes four consecutive k of a row (A) or column (B) per register. Four consecutive
|
||||
// weights are one code byte; four codes of one register instead come from the same 2-bit field of four consecutive
|
||||
// bytes - (word >> 2t) & 0x03030303 holds weights t, t+4, t+8, t+12 of a 16-weight word. So the kernels use that order
|
||||
// for the k of a half-block: mma k index 16h + 4t + j stands for weight (and activation) 16h + t + 4j. The
|
||||
// activations are stored in it (perm32), and a lane's two B registers (h = 0 and 1) are adjacent: one 8-byte load.
|
||||
#include "strata/prefill/moe_fused.hpp"
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <mutex>
|
||||
|
||||
namespace strata::prefill::fused {
|
||||
namespace {
|
||||
|
||||
void ck(cudaError_t e, const char* what) {
|
||||
if (e != cudaSuccess) {
|
||||
std::fprintf(stderr, "prefill fused experts: %s: %s\n", what, cudaGetErrorString(e));
|
||||
std::exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
constexpr int AB = 80; // activation bytes per 64 values: 64 codes, then {d0, c0, d1, c1}
|
||||
constexpr int MAGIC = 0x4B400000; // the bits of 1.5 * 2^23
|
||||
constexpr float MAGICF = 12582912.0f;
|
||||
// The Strata Q2_0 blob (moe_mmq.cu strata_q2_kernel, kernels.cu blob_dequant_kernel): gate/up codes [1280][640 B]
|
||||
// (rows interleaved: gate of feature f at row 2f, its up at 2f + 1), down codes [2560][160 B], gate/up scales
|
||||
// [1280][40] fp16, down scales [2560][10] fp16.
|
||||
constexpr size_t O_GU_CODES = 0, O_D_CODES = (size_t) 1280 * 640, O_GU_SC = O_D_CODES + (size_t) 2560 * 160,
|
||||
O_D_SC = O_GU_SC + (size_t) 1280 * 40 * 2;
|
||||
constexpr int THREADS = 512; // 16 warps: 4 along the weight rows (64 each) x 4 along the tile rows (16 each)
|
||||
constexpr int WROWS = 256; // weight rows per work item: 128 features of gate/up, or 256 outputs of down
|
||||
constexpr int STAGES = 4; // cp.async pipeline depth, one 64-weight block of K per stage
|
||||
constexpr int STAGE_BYTES = WROWS * 16 + kTileRows * AB;
|
||||
constexpr size_t smem_bytes(bool gu) {
|
||||
return (size_t) WROWS * (gu ? 40 : 10) * 2 + STAGES * STAGE_BYTES + kTileRows * 4;
|
||||
}
|
||||
|
||||
// the stored position of value k of a 32-value half-block (see "The K order")
|
||||
__host__ __device__ __forceinline__ int perm32(int k) { return 8 * (k & 3) + 4 * (k >> 4) + ((k & 15) >> 2); }
|
||||
|
||||
struct Tables {
|
||||
int32_t *cnt, *off, *fill, *ts; // per expert: rows, first row (E + 1), placement cursor, first tile (E + 1)
|
||||
int2* tiles; // {expert, first row}, in expert order
|
||||
};
|
||||
size_t tiles_at(int n_expert) { return ((size_t) (4 * n_expert + 2) * 4 + 15) / 16 * 16; }
|
||||
Tables tables(void* scratch, int n_expert) {
|
||||
int32_t* p = (int32_t*) scratch;
|
||||
return {p, p + n_expert, p + 2 * n_expert + 1, p + 3 * n_expert + 1,
|
||||
(int2*) ((uint8_t*) scratch + tiles_at(n_expert))};
|
||||
}
|
||||
|
||||
// ---- activations: one warp per 64 values
|
||||
__global__ void quant_act_kernel(const float* __restrict__ x, int64_t nblk, uint8_t* __restrict__ xa) {
|
||||
const int64_t w = (int64_t) blockIdx.x * (blockDim.x / 32) + threadIdx.x / 32;
|
||||
const int lane = threadIdx.x & 31;
|
||||
if (w >= nblk) return;
|
||||
const float v0 = x[w * 64 + lane], v1 = x[w * 64 + 32 + lane];
|
||||
float a0 = fabsf(v0), a1 = fabsf(v1);
|
||||
#pragma unroll
|
||||
for (int o = 16; o > 0; o >>= 1) {
|
||||
a0 = fmaxf(a0, __shfl_xor_sync(0xffffffffu, a0, o));
|
||||
a1 = fmaxf(a1, __shfl_xor_sync(0xffffffffu, a1, o));
|
||||
}
|
||||
const int q0 = a0 > 0.0f ? __float2int_rn(v0 * (127.0f / a0)) : 0;
|
||||
const int q1 = a1 > 0.0f ? __float2int_rn(v1 * (127.0f / a1)) : 0;
|
||||
int s0 = q0, s1 = q1;
|
||||
#pragma unroll
|
||||
for (int o = 16; o > 0; o >>= 1) {
|
||||
s0 += __shfl_xor_sync(0xffffffffu, s0, o);
|
||||
s1 += __shfl_xor_sync(0xffffffffu, s1, o);
|
||||
}
|
||||
uint8_t* out = xa + w * AB;
|
||||
out[perm32(lane)] = (uint8_t) (int8_t) q0;
|
||||
out[32 + perm32(lane)] = (uint8_t) (int8_t) q1;
|
||||
if (lane == 0)
|
||||
*(float4*) (out + 64) = make_float4(a0 / 127.0f, -(MAGICF + (float) s0), a1 / 127.0f, -(MAGICF + (float) s1));
|
||||
}
|
||||
|
||||
// ---- grouping: counts, offsets and tiles, placement
|
||||
__global__ void count_kernel(const int32_t* __restrict__ ids, int64_t n, int E, int32_t* __restrict__ cnt) {
|
||||
extern __shared__ int32_t hist[];
|
||||
for (int e = threadIdx.x; e < E; e += blockDim.x) hist[e] = 0;
|
||||
__syncthreads();
|
||||
for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x * blockDim.x) {
|
||||
const int e = ids[i];
|
||||
if ((unsigned) e < (unsigned) E) atomicAdd(&hist[e], 1);
|
||||
}
|
||||
__syncthreads();
|
||||
for (int e = threadIdx.x; e < E; e += blockDim.x)
|
||||
if (hist[e]) atomicAdd(&cnt[e], hist[e]);
|
||||
}
|
||||
|
||||
// exclusive prefix sum over the block (blockDim a multiple of 32); *total = the block's sum
|
||||
__device__ int block_excl_sum(int v, int* total, int* sh) {
|
||||
const int lane = threadIdx.x & 31, w = threadIdx.x >> 5, nw = blockDim.x >> 5;
|
||||
int x = v;
|
||||
#pragma unroll
|
||||
for (int o = 1; o < 32; o <<= 1) {
|
||||
const int y = __shfl_up_sync(0xffffffffu, x, o);
|
||||
if (lane >= o) x += y;
|
||||
}
|
||||
if (lane == 31) sh[w] = x;
|
||||
__syncthreads();
|
||||
if (w == 0) {
|
||||
int s = lane < nw ? sh[lane] : 0;
|
||||
#pragma unroll
|
||||
for (int o = 1; o < 32; o <<= 1) {
|
||||
const int y = __shfl_up_sync(0xffffffffu, s, o);
|
||||
if (lane >= o) s += y;
|
||||
}
|
||||
sh[lane] = s;
|
||||
}
|
||||
__syncthreads();
|
||||
const int r = x - v + (w > 0 ? sh[w - 1] : 0);
|
||||
*total = sh[nw - 1];
|
||||
__syncthreads(); // sh is reused by the next call
|
||||
return r;
|
||||
}
|
||||
|
||||
// one block, a thread per expert (E <= 1024)
|
||||
__global__ void scan_kernel(Tables tb, int E) {
|
||||
__shared__ int sh[32];
|
||||
const int e = threadIdx.x;
|
||||
const int c = e < E ? tb.cnt[e] : 0, nt = (c + kTileRows - 1) / kTileRows;
|
||||
int rows = 0, tiles = 0;
|
||||
const int o = block_excl_sum(c, &rows, sh), to = block_excl_sum(nt, &tiles, sh);
|
||||
if (e < E) {
|
||||
tb.off[e] = o;
|
||||
tb.fill[e] = o;
|
||||
tb.ts[e] = to;
|
||||
for (int i = 0; i < nt; ++i) tb.tiles[to + i] = make_int2(e, o + i * kTileRows);
|
||||
}
|
||||
if (e == 0) {
|
||||
tb.off[E] = rows;
|
||||
tb.ts[E] = tiles;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void place_kernel(const int32_t* __restrict__ ids, int64_t n, int k, int E, Tables tb,
|
||||
int32_t* __restrict__ slot, int32_t* __restrict__ src) {
|
||||
const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= n) return;
|
||||
const int e = ids[i];
|
||||
if ((unsigned) e >= (unsigned) E) { slot[i] = 0; return; } // (the router never writes one)
|
||||
const int p = atomicAdd(&tb.fill[e], 1);
|
||||
slot[i] = p;
|
||||
src[p] = (int32_t) (i / k);
|
||||
}
|
||||
|
||||
// ---- the expert products
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
__device__ __forceinline__ void cp16(void* dst, const void* src) {
|
||||
asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"((unsigned) __cvta_generic_to_shared(dst)),
|
||||
"l"(src));
|
||||
}
|
||||
__device__ __forceinline__ void cp4(void* dst, const void* src) {
|
||||
asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" ::"r"((unsigned) __cvta_generic_to_shared(dst)),
|
||||
"l"(src));
|
||||
}
|
||||
__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;\n" ::); }
|
||||
template <int N> __device__ __forceinline__ void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); }
|
||||
// d = A (16 x 32 s8, row) * B (32 x 8 s8, col), int32, from zero
|
||||
__device__ __forceinline__ void mma_s8(int (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) {
|
||||
asm("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%10,%10,%10};\n"
|
||||
: "=r"(d[0]), "=r"(d[1]), "=r"(d[2]), "=r"(d[3])
|
||||
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1), "r"(0));
|
||||
}
|
||||
#endif
|
||||
|
||||
// One work item = (tile of up to 64 routed rows of expert e, a block of 256 weight rows); a persistent grid walks
|
||||
// the batch's items (their count is on the device). GU: gate/up rows of 128 features against the rows' tokens'
|
||||
// activations (`act` per token, 2560 values), SwiGLU, H to int8 per 32 features into `out` (per row, 640 values).
|
||||
// !GU: down rows (256 outputs) against H (`act` per row), FP32 into `dm` per row.
|
||||
// Warp (wf, wt): weight rows 64 wf.. as four m16 tiles (GU: a tile = 8 features, its gate rows as mma rows 0-7 and
|
||||
// its up rows as 8-15, so a lane holds gate and up of one feature: SwiGLU in registers), tile rows 16 wt.. as two n8.
|
||||
template <bool GU>
|
||||
__global__ void __launch_bounds__(THREADS, 1)
|
||||
expert_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ act, const int32_t* __restrict__ src,
|
||||
uint8_t* __restrict__ out, float* __restrict__ dm) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
constexpr int NKB = GU ? 40 : 10; // 64-weight blocks along K
|
||||
constexpr int NFB = GU ? 5 : 10; // weight-row blocks per tile
|
||||
constexpr int CODE_LD = GU ? 640 : 160; // code bytes per weight row
|
||||
constexpr size_t O_CODES = GU ? O_GU_CODES : O_D_CODES, O_SC = GU ? O_GU_SC : O_D_SC;
|
||||
constexpr int ACT_LD = NKB * AB; // bytes per activation row: 3200 (a token), 800 (a row's H)
|
||||
constexpr unsigned M2 = 0x03030303u;
|
||||
extern __shared__ __align__(16) uint8_t smem[];
|
||||
const __half* wsc = (const __half*) smem; // [256][NKB] the work item's weight scales
|
||||
uint8_t* stages = smem + WROWS * NKB * 2; // [STAGES] x {codes [256][16], act [64][AB]}
|
||||
int* srow = (int*) (stages + STAGES * STAGE_BYTES); // [64] the activation row of each tile row
|
||||
|
||||
const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, tig = lane & 3;
|
||||
const int wf = warp & 3, nb0 = 16 * (warp >> 2);
|
||||
int ra[4], rb[4]; // the local weight rows of mma rows g and g + 8
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
if (GU) { ra[i] = 2 * (32 * wf + 8 * i + g); rb[i] = ra[i] + 1; }
|
||||
else { ra[i] = 64 * wf + 16 * i + g; rb[i] = ra[i] + 8; }
|
||||
}
|
||||
const int t0 = tb.ts[b.e0], nwork = (tb.ts[b.e1] - t0) * NFB;
|
||||
for (int w = blockIdx.x; w < nwork; w += gridDim.x) {
|
||||
const int2 tl = tb.tiles[t0 + w / NFB];
|
||||
const int fb = w % NFB, e = tl.x, row0 = tl.y, nrows = min(kTileRows, tb.off[e + 1] - row0);
|
||||
const uint8_t* blob = b.blob[e - b.e0];
|
||||
const int rbase = fb * WROWS;
|
||||
__syncthreads(); // the previous item is done with the buffers
|
||||
if (tid < kTileRows) {
|
||||
const int r = row0 + min(tid, nrows - 1); // rows past the tile's end repeat its last one
|
||||
srow[tid] = GU ? src[r] : r;
|
||||
}
|
||||
for (int c = tid; c < WROWS * 5; c += THREADS) { // the scales: 5 x 16 B (GU) or 5 x 4 B a row
|
||||
const int r = c / 5, q = c % 5;
|
||||
if (GU) cp16(smem + r * 80 + q * 16, blob + O_SC + (size_t) (rbase + r) * 80 + q * 16);
|
||||
else cp4(smem + r * 20 + q * 4, blob + O_SC + (size_t) (rbase + r) * 20 + q * 4);
|
||||
}
|
||||
__syncthreads(); // srow
|
||||
auto load = [&](int kb) {
|
||||
uint8_t* st = stages + (kb % STAGES) * STAGE_BYTES;
|
||||
if (tid < WROWS) cp16(st + tid * 16, blob + O_CODES + (size_t) (rbase + tid) * CODE_LD + kb * 16);
|
||||
const int a = tid - (THREADS - kTileRows * 5); // the last 320 threads: 64 rows x 5 x 16 B
|
||||
if (a >= 0) {
|
||||
const int r = a / 5, q = a % 5;
|
||||
cp16(st + WROWS * 16 + r * AB + q * 16, act + (size_t) srow[r] * ACT_LD + kb * AB + q * 16);
|
||||
}
|
||||
};
|
||||
#pragma unroll
|
||||
for (int s = 0; s < STAGES - 1; ++s) {
|
||||
if (s < NKB) load(s);
|
||||
cp_commit();
|
||||
}
|
||||
// a warp whose rows are all past the tile's end only takes part in the loads
|
||||
const bool on0 = nb0 < nrows, on1 = nb0 + 8 < nrows;
|
||||
float acc[4][2][4];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i)
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 2; ++n)
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 4; ++q) acc[i][n][q] = 0.0f;
|
||||
for (int kb = 0; kb < NKB; ++kb) {
|
||||
cp_wait<STAGES - 2>();
|
||||
__syncthreads();
|
||||
if (kb + STAGES - 1 < NKB) load(kb + STAGES - 1);
|
||||
cp_commit();
|
||||
if (!on0) continue;
|
||||
const uint8_t* st = stages + (kb % STAGES) * STAGE_BYTES;
|
||||
const uint8_t* sa = st + WROWS * 16;
|
||||
uint2 bq[2][2]; // [n8 tile][half]: the B registers
|
||||
float4 sc[2][2]; // [n8 tile][column 2 tig + c]: {d0, c0, d1, c1}
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 2; ++n) {
|
||||
if (n == 1 && !on1) break;
|
||||
const uint8_t* r = sa + (nb0 + 8 * n + g) * AB + 8 * tig;
|
||||
bq[n][0] = *(const uint2*) r;
|
||||
bq[n][1] = *(const uint2*) (r + 32);
|
||||
sc[n][0] = *(const float4*) (sa + (nb0 + 8 * n + 2 * tig) * AB + 64);
|
||||
sc[n][1] = *(const float4*) (sa + (nb0 + 8 * n + 2 * tig + 1) * AB + 64);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
const uint4 ca = *(const uint4*) (st + ra[i] * 16), cb = *(const uint4*) (st + rb[i] * 16);
|
||||
const float da = __half2float(wsc[ra[i] * NKB + kb]), db = __half2float(wsc[rb[i] * NKB + kb]);
|
||||
const int s = 2 * tig;
|
||||
const uint32_t a0[4] = {(ca.x >> s) & M2, (cb.x >> s) & M2, (ca.y >> s) & M2, (cb.y >> s) & M2};
|
||||
const uint32_t a1[4] = {(ca.z >> s) & M2, (cb.z >> s) & M2, (ca.w >> s) & M2, (cb.w >> s) & M2};
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 2; ++n) {
|
||||
if (n == 1 && !on1) break;
|
||||
int d0[4], d1[4];
|
||||
mma_s8(d0, a0, bq[n][0].x, bq[n][0].y);
|
||||
mma_s8(d1, a1, bq[n][1].x, bq[n][1].y);
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 4; ++q) {
|
||||
const float4 z = sc[n][q & 1];
|
||||
const float u0 = __int_as_float(MAGIC + d0[q]) + z.y, u1 = __int_as_float(MAGIC + d1[q]) + z.w;
|
||||
acc[i][n][q] = fmaf(q < 2 ? da : db, fmaf(z.z, u1, z.x * u0), acc[i][n][q]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!on0) continue;
|
||||
if (GU) {
|
||||
// SwiGLU of feature 32 wf + 8 i + g (lane g of m tile i), tile rows nb0 + 8 n + 2 tig + c; H to int8 per
|
||||
// row over the warp's 32 features (= one 32-value half-block of the down product's K)
|
||||
const int blk = 4 * fb + wf;
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 2; ++n) {
|
||||
if (n == 1 && !on1) break;
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 2; ++c) {
|
||||
float h[4], am = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
const float gt = acc[i][n][c], up = acc[i][n][2 + c];
|
||||
h[i] = gt / (1.0f + __expf(-gt)) * up;
|
||||
am = fmaxf(am, fabsf(h[i]));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int o = 4; o < 32; o <<= 1) am = fmaxf(am, __shfl_xor_sync(0xffffffffu, am, o));
|
||||
const float inv = am > 0.0f ? 127.0f / am : 0.0f;
|
||||
int qv[4], sum = 0;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) { qv[i] = __float2int_rn(h[i] * inv); sum += qv[i]; }
|
||||
#pragma unroll
|
||||
for (int o = 4; o < 32; o <<= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o);
|
||||
const int r = nb0 + 8 * n + 2 * tig + c;
|
||||
if (r < nrows) {
|
||||
uint8_t* o = out + (size_t) (row0 + r) * (10 * AB) + (blk >> 1) * AB;
|
||||
const int hh = blk & 1;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) o[32 * hh + perm32(8 * i + g)] = (uint8_t) (int8_t) qv[i];
|
||||
if (g == 0) *(float2*) (o + 64 + 8 * hh) = make_float2(am / 127.0f, -(MAGICF + (float) sum));
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i)
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 2; ++n) {
|
||||
if (n == 1 && !on1) break;
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 4; ++q) {
|
||||
const int r = nb0 + 8 * n + 2 * tig + (q & 1);
|
||||
if (r < nrows) dm[(size_t) (row0 + r) * 2560 + rbase + (q < 2 ? ra[i] : rb[i])] = acc[i][n][q];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
unsigned blocks(int64_t n, int per) { return (unsigned) ((n + per - 1) / per); }
|
||||
|
||||
struct DevInfo {
|
||||
bool done = false;
|
||||
int cc = 0, sms = 0, occ_gu = 0, occ_d = 0;
|
||||
};
|
||||
std::mutex g_mu;
|
||||
DevInfo g_dev[32];
|
||||
|
||||
const DevInfo& dev_info() {
|
||||
int dev = 0;
|
||||
cudaGetDevice(&dev);
|
||||
std::lock_guard<std::mutex> lk(g_mu);
|
||||
DevInfo& d = g_dev[dev & 31];
|
||||
if (d.done) return d;
|
||||
d.done = true;
|
||||
int major = 0, minor = 0;
|
||||
cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev);
|
||||
cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev);
|
||||
cudaDeviceGetAttribute(&d.sms, cudaDevAttrMultiProcessorCount, dev);
|
||||
d.cc = 10 * major + minor;
|
||||
if (d.cc < 80) return d;
|
||||
// the device code of this card's arch must have the kernels (a build without an sm_80+ target has an empty body)
|
||||
cudaFuncAttributes fa{};
|
||||
if (cudaFuncGetAttributes(&fa, expert_kernel<true>) != cudaSuccess || fa.ptxVersion < 80) {
|
||||
cudaGetLastError();
|
||||
d.cc = 0;
|
||||
return d;
|
||||
}
|
||||
ck(cudaFuncSetAttribute(expert_kernel<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int) smem_bytes(true)),
|
||||
"smem gu");
|
||||
ck(cudaFuncSetAttribute(expert_kernel<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int) smem_bytes(false)),
|
||||
"smem down");
|
||||
cudaOccupancyMaxActiveBlocksPerMultiprocessor(&d.occ_gu, expert_kernel<true>, THREADS, smem_bytes(true));
|
||||
cudaOccupancyMaxActiveBlocksPerMultiprocessor(&d.occ_d, expert_kernel<false>, THREADS, smem_bytes(false));
|
||||
if (d.occ_gu < 1 || d.occ_d < 1) d.cc = 0; // does not fit this card: the MMQ path
|
||||
cudaGetLastError();
|
||||
return d;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
bool built() { return true; }
|
||||
bool available() { return dev_info().cc >= 80; }
|
||||
// 0.1.36: on by default for the Q2_0 pack (+13-23% prompts on an RTX 5070; teacher-forced against the FP16 path as
|
||||
// close as MMQ at 8K and closer at 32K: phaseA-tests s18-tf); STRATA_PF_FUSED=0 keeps MMQ. The native packs' kernels
|
||||
// (moe_fused_iq) stay opt-in: STRATA_PF_FUSED=1 (requested()).
|
||||
bool enabled() {
|
||||
static const bool env = [] {
|
||||
const char* v = std::getenv("STRATA_PF_FUSED");
|
||||
return v == nullptr || v[0] != '0';
|
||||
}();
|
||||
return env && available();
|
||||
}
|
||||
bool requested() {
|
||||
static const bool env = [] {
|
||||
const char* v = std::getenv("STRATA_PF_FUSED");
|
||||
return v != nullptr && v[0] == '1';
|
||||
}();
|
||||
return env && available();
|
||||
}
|
||||
|
||||
size_t act_bytes(int64_t rows, int64_t cols) { return (size_t) rows * (size_t) (cols / 64) * AB; }
|
||||
size_t group_bytes(int64_t n, int n_expert) {
|
||||
return tiles_at(n_expert) + ((size_t) (n + kTileRows - 1) / kTileRows + (size_t) n_expert) * sizeof(int2);
|
||||
}
|
||||
|
||||
void quantize_act(const float* x, int64_t rows, int64_t cols, void* xa, void* stream) {
|
||||
if (rows <= 0) return;
|
||||
const int64_t nblk = rows * (cols / 64);
|
||||
quant_act_kernel<<<blocks(nblk, 8), 256, 0, (cudaStream_t) stream>>>(x, nblk, (uint8_t*) xa);
|
||||
ck(cudaGetLastError(), "quantize_act");
|
||||
}
|
||||
|
||||
void group(const int32_t* ids, int64_t n, int k, int n_expert, void* scratch, int32_t* slot, int32_t* src,
|
||||
void* stream) {
|
||||
if (n_expert < 1 || n_expert > 1024) {
|
||||
std::fprintf(stderr, "prefill fused experts: %d experts (the grouping takes 1 to 1024)\n", n_expert);
|
||||
std::exit(1);
|
||||
}
|
||||
const cudaStream_t s = (cudaStream_t) stream;
|
||||
const Tables tb = tables(scratch, n_expert);
|
||||
ck(cudaMemsetAsync(tb.cnt, 0, (size_t) n_expert * 4, s), "group memset");
|
||||
if (n > 0) {
|
||||
const unsigned nb = std::min<unsigned>(blocks(n, 1024), 4u * (unsigned) std::max(dev_info().sms, 1));
|
||||
count_kernel<<<nb, 1024, (size_t) n_expert * 4, s>>>(ids, n, n_expert, tb.cnt);
|
||||
}
|
||||
scan_kernel<<<1, (unsigned) ((n_expert + 31) / 32 * 32), 0, s>>>(tb, n_expert);
|
||||
if (n > 0) place_kernel<<<blocks(n, 256), 256, 0, s>>>(ids, n, k, n_expert, tb, slot, src);
|
||||
ck(cudaGetLastError(), "group");
|
||||
}
|
||||
|
||||
void experts(const Batch& b, int n_expert, int64_t n, const void* scratch, const void* xa, const int32_t* src,
|
||||
void* ha, float* dm, void* stream) {
|
||||
if (b.e1 <= b.e0 || n <= 0) return;
|
||||
const DevInfo& d = dev_info();
|
||||
const cudaStream_t s = (cudaStream_t) stream;
|
||||
const Tables tb = tables(const_cast<void*>(scratch), n_expert);
|
||||
// the most tiles the batch can have: every row in it, plus a partial tile per expert
|
||||
const int64_t tiles = (n + kTileRows - 1) / kTileRows + (b.e1 - b.e0);
|
||||
const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * 5, (int64_t) d.sms * d.occ_gu);
|
||||
const unsigned g_d = (unsigned) std::min<int64_t>(tiles * 10, (int64_t) d.sms * d.occ_d);
|
||||
expert_kernel<true><<<g_gu, THREADS, smem_bytes(true), s>>>(b, tb, (const uint8_t*) xa, src, (uint8_t*) ha,
|
||||
nullptr);
|
||||
expert_kernel<false><<<g_d, THREADS, smem_bytes(false), s>>>(b, tb, (const uint8_t*) ha, src, nullptr, dm);
|
||||
ck(cudaGetLastError(), "experts");
|
||||
}
|
||||
|
||||
} // namespace strata::prefill::fused
|
||||
@@ -0,0 +1,630 @@
|
||||
// src/prefill/moe_fused_iq.cu - see include/strata/prefill/moe_fused_iq.hpp (#136: the native packs' prompt experts on
|
||||
// the fused int8 kernels of moe_fused.cu).
|
||||
//
|
||||
// The arithmetic. An activation block of 32 values is x = d_x * a, a = round(x / d_x) in -127..127, stored in natural
|
||||
// order (per 64 values: 64 codes, then {d_0, d_1} and 8 unused bytes - 80 bytes, the Q2_0 path's size). A weight
|
||||
// block of 32 (or 16) values is w = d_w * q with q an int8: the i-quants' codebook entries with their signs applied
|
||||
// (IQ2_XXS / IQ2_XS / IQ2_S: 0..43, IQ3_XXS: 0..62, IQ3_S: 0..15, IQ4_XS / IQ4_NL: -127..113), Q2_0's code - 1. So
|
||||
// sum w x = d_w * d_x * sum q a, |sum q a| <= 32 * 127 * 127 < 2^22,
|
||||
// and the int32 dot becomes a float exactly with one float add: the mma starts from C = 0x4B400000 (the bits of
|
||||
// 1.5 * 2^23), so as_float(D) - 1.5 * 2^23 is the dot. Per output and 32 values: that add, d_w * d_x, one fma (the
|
||||
// formats with a scale per 16 values: two dots, each times its d_w, then d_x). d_w is llama.cpp's own: IQ2_XXS
|
||||
// d (2s + 1) / 8, IQ2_XS / IQ2_S d (2s + 1) / 8 per 16, IQ3_XXS d (2s + 1) / 4, IQ3_S d (2s + 1), IQ4_XS d (s - 32),
|
||||
// IQ4_NL and Q2_0 d (dequantize_row_* in ggml-quants.c, mmq-load-tiles.cuh).
|
||||
//
|
||||
// The load stage. Per 64 values of K (a stage), a thread decodes one 32-value sub-block of one of the work item's
|
||||
// weight rows: its block bytes are read from the blob with 16-bit loads (the GGUF blocks are 2-byte aligned: 66 / 74 /
|
||||
// 82 / 98 / 110 / 136 / 18 bytes) one stage ahead into registers (and two 256-value blocks ahead into L2), then turned
|
||||
// into 32 int8 and the sub-block's scales in shared memory, double-buffered. The codebooks live in shared memory (up
|
||||
// to 8 KB, IQ2_S). The activations arrive by cp.async, four stages deep, as in moe_fused.cu; the fragments come by
|
||||
// ldmatrix. The products are mma.sync m16n8k32 (m16n8k16 for the formats with a scale per 16 values); the epilogues
|
||||
// (SwiGLU, H to int8 per 32 features; down into the per-slot rows) are moe_fused.cu's.
|
||||
#include "strata/prefill/moe_fused_iq.hpp"
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include "ggml.h"
|
||||
|
||||
#define GGML_COMMON_DECL_CUDA
|
||||
#define GGML_COMMON_IMPL_CUDA
|
||||
#include "ggml-common.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <mutex>
|
||||
|
||||
namespace strata::prefill::fused {
|
||||
namespace {
|
||||
|
||||
void ck(cudaError_t e, const char* what) {
|
||||
if (e != cudaSuccess) {
|
||||
std::fprintf(stderr, "prefill fused experts (native): %s: %s\n", what, cudaGetErrorString(e));
|
||||
std::exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
constexpr int AB = 80; // activation bytes per 64 values: 64 codes, {d0, d1}, 8 unused
|
||||
constexpr int MAGIC = 0x4B400000; // the bits of 1.5 * 2^23
|
||||
constexpr float MAGICF = 12582912.0f;
|
||||
// 16 warps, each 64 weight rows x 16 routed rows; WW of them along the weight rows and 16 / WW along the routed rows: a
|
||||
// work item is (WR = 64 WW weight rows) x (TR = 256 / WW routed rows of one expert) - 256 x 64 or 128 x 128, chosen
|
||||
// per launch by the rows per expert (pick_ww)
|
||||
constexpr int THREADS = 512;
|
||||
constexpr int WLD = 80; // bytes per decoded weight row of a stage: 64 int8 + 16 (no bank conflicts)
|
||||
__host__ __device__ constexpr int tile_rows(int ww) { return 256 / ww; }
|
||||
__host__ __device__ constexpr int weight_rows(int ww) { return 64 * ww; }
|
||||
constexpr int ASTAGES = 4; // cp.async depth of the activations (64 values of K a stage)
|
||||
constexpr int GU_ROWS_K = 2560, D_ROWS_K = 640; // K of gate/up (n_embd) and of down (n_ff)
|
||||
|
||||
// the formats (ggml type ids)
|
||||
constexpr int T_IQ2_XXS = GGML_TYPE_IQ2_XXS, T_IQ2_XS = GGML_TYPE_IQ2_XS, T_IQ2_S = GGML_TYPE_IQ2_S,
|
||||
T_IQ3_XXS = GGML_TYPE_IQ3_XXS, T_IQ3_S = GGML_TYPE_IQ3_S, T_IQ4_XS = GGML_TYPE_IQ4_XS,
|
||||
T_IQ4_NL = GGML_TYPE_IQ4_NL, T_Q2_0 = GGML_TYPE_Q2_0;
|
||||
static_assert(sizeof(block_iq2_xxs) == 66 && sizeof(block_iq2_xs) == 74 && sizeof(block_iq2_s) == 82 &&
|
||||
sizeof(block_iq3_xxs) == 98 && sizeof(block_iq3_s) == 110 && sizeof(block_iq4_xs) == 136 &&
|
||||
sizeof(block_iq4_nl) == 18 && sizeof(block_q2_0) == 18, "the block layouts this file decodes");
|
||||
|
||||
// block bytes, scale per 16 values, codebook bytes in shared memory
|
||||
__host__ __device__ constexpr int block_bytes(int t) {
|
||||
return t == T_IQ2_XXS ? 66 : t == T_IQ2_XS ? 74 : t == T_IQ2_S ? 82 : t == T_IQ3_XXS ? 98 : t == T_IQ3_S ? 110
|
||||
: t == T_IQ4_XS ? 136 : 18;
|
||||
}
|
||||
__host__ __device__ constexpr bool per16(int t) { return t == T_IQ2_XS || t == T_IQ2_S; }
|
||||
__host__ __device__ constexpr int grid_bytes(int t) {
|
||||
return t == T_IQ2_XXS ? 256 * 8 : t == T_IQ2_XS ? 512 * 8 : t == T_IQ2_S ? 1024 * 8 : t == T_IQ3_XXS ? 256 * 4
|
||||
: t == T_IQ3_S ? 512 * 4 : 0;
|
||||
}
|
||||
__host__ __device__ constexpr size_t smem_bytes(int t, int ww) {
|
||||
return (size_t) 2 * weight_rows(ww) * (WLD + 16) + (size_t) ASTAGES * tile_rows(ww) * AB +
|
||||
(size_t) tile_rows(ww) * 4 + grid_bytes(t);
|
||||
}
|
||||
|
||||
// moe_fused.cu's grouping tables (the same layout: fused::group writes them)
|
||||
struct Tables {
|
||||
int32_t *cnt, *off, *fill, *ts;
|
||||
int2* tiles;
|
||||
};
|
||||
size_t tiles_at(int n_expert) { return ((size_t) (4 * n_expert + 2) * 4 + 15) / 16 * 16; }
|
||||
Tables tables(void* scratch, int n_expert) {
|
||||
int32_t* p = (int32_t*) scratch;
|
||||
return {p, p + n_expert, p + 2 * n_expert + 1, p + 3 * n_expert + 1,
|
||||
(int2*) ((uint8_t*) scratch + tiles_at(n_expert))};
|
||||
}
|
||||
|
||||
// ---- activations in natural order: one warp per 64 values
|
||||
__global__ void quant_act_nat_kernel(const float* __restrict__ x, int64_t nblk, uint8_t* __restrict__ xa) {
|
||||
const int64_t w = (int64_t) blockIdx.x * (blockDim.x / 32) + threadIdx.x / 32;
|
||||
const int lane = threadIdx.x & 31;
|
||||
if (w >= nblk) return;
|
||||
const float v0 = x[w * 64 + lane], v1 = x[w * 64 + 32 + lane];
|
||||
float a0 = fabsf(v0), a1 = fabsf(v1);
|
||||
#pragma unroll
|
||||
for (int o = 16; o > 0; o >>= 1) {
|
||||
a0 = fmaxf(a0, __shfl_xor_sync(0xffffffffu, a0, o));
|
||||
a1 = fmaxf(a1, __shfl_xor_sync(0xffffffffu, a1, o));
|
||||
}
|
||||
uint8_t* out = xa + w * AB;
|
||||
out[lane] = (uint8_t) (int8_t) (a0 > 0.0f ? __float2int_rn(v0 * (127.0f / a0)) : 0);
|
||||
out[32 + lane] = (uint8_t) (int8_t) (a1 > 0.0f ? __float2int_rn(v1 * (127.0f / a1)) : 0);
|
||||
if (lane == 0) *(float4*) (out + 64) = make_float4(a0 / 127.0f, a1 / 127.0f, 0.0f, 0.0f);
|
||||
}
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
__device__ __forceinline__ void cp16(void* dst, const void* src) {
|
||||
asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"((unsigned) __cvta_generic_to_shared(dst)),
|
||||
"l"(src));
|
||||
}
|
||||
__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;\n" ::); }
|
||||
__device__ __forceinline__ void pf_l2(const void* p) { asm volatile("prefetch.global.L2 [%0];\n" ::"l"(p)); }
|
||||
template <int N> __device__ __forceinline__ void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); }
|
||||
// d = MAGIC + A (16 x 32 s8, row) * B (32 x 8 s8, col): the int32 dot with the magic bias already added (the C
|
||||
// operand), so as_float(d) - 1.5 * 2^23 is the dot as a float
|
||||
__device__ __forceinline__ void mma32(int (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) {
|
||||
asm("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%10,%10,%10};\n"
|
||||
: "=r"(d[0]), "=r"(d[1]), "=r"(d[2]), "=r"(d[3])
|
||||
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1), "r"(MAGIC));
|
||||
}
|
||||
// the same at m16n8k16 (A 16 x 16, B 16 x 8)
|
||||
__device__ __forceinline__ void mma16(int (&d)[4], uint32_t a0, uint32_t a1, uint32_t b0) {
|
||||
asm("mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%7,%7,%7};\n"
|
||||
: "=r"(d[0]), "=r"(d[1]), "=r"(d[2]), "=r"(d[3])
|
||||
: "r"(a0), "r"(a1), "r"(b0), "r"(MAGIC));
|
||||
}
|
||||
__device__ __forceinline__ float dotf(int d) { return __int_as_float(d) - MAGICF; }
|
||||
// four 8 x 16-byte matrices, row addresses from lanes 0-7, 8-15, 16-23, 24-31: lane (g, t) gets bytes 4t..4t+3 of row g
|
||||
// of each - an mma fragment
|
||||
__device__ __forceinline__ void ldsm4(uint32_t (&r)[4], const void* p) {
|
||||
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
|
||||
: "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
|
||||
: "r"((unsigned) __cvta_generic_to_shared(p))
|
||||
: "memory");
|
||||
}
|
||||
|
||||
// ---- the load stage: a 32-value sub-block's bytes into registers, then int8 and scales
|
||||
__device__ __forceinline__ uint32_t ld16(const uint8_t* p) { return *(const uint16_t*) p; }
|
||||
__device__ __forceinline__ uint32_t ld32(const uint8_t* p) { return ld16(p) | (ld16(p + 2) << 16); }
|
||||
__device__ __forceinline__ float half_at(uint32_t w) { return __half2float(__ushort_as_half((unsigned short) w)); }
|
||||
// llama.cpp's sign unpacking: 7 bits of signs, the 8th their parity (bit 7 of v may be anything)
|
||||
__device__ __forceinline__ uint32_t unpack_ksigns(uint32_t v) {
|
||||
v &= 0xFF;
|
||||
const uint32_t p = __popc(v) & 1;
|
||||
return (v ^ p << 7) * 0x01010101u;
|
||||
}
|
||||
// 8 bytes of a codebook entry with the sign byte `s` (broadcast) applied: bits 0-3 to .x, 4-7 to .y
|
||||
__device__ __forceinline__ void signed8(uint32_t gx, uint32_t gy, uint32_t s, uint32_t& qx, uint32_t& qy) {
|
||||
const uint32_t m0 = __vcmpne4(s & 0x08040201u, 0), m1 = __vcmpne4(s & 0x80402010u, 0);
|
||||
qx = __vsub4(gx ^ m0, m0);
|
||||
qy = __vsub4(gy ^ m1, m1);
|
||||
}
|
||||
// 8 nibbles of q4 through a 16-entry int8 table (4 words): the low nibbles' values in .x, the high ones' in .y
|
||||
__device__ __forceinline__ void table16(uint32_t q4, const uint32_t (&t)[4], uint32_t& lo, uint32_t& hi) {
|
||||
uint32_t tmp[2];
|
||||
const uint32_t sel = 0x32103210u | ((q4 & 0x88888888u) >> 1);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 2; ++i) {
|
||||
const uint32_t sh = 16 * i;
|
||||
const uint32_t l = __byte_perm(t[0], t[1], q4 >> sh), h = __byte_perm(t[2], t[3], q4 >> sh);
|
||||
tmp[i] = __byte_perm(l, h, sel >> sh);
|
||||
}
|
||||
lo = __byte_perm(tmp[0], tmp[1], 0x6420);
|
||||
hi = __byte_perm(tmp[0], tmp[1], 0x7531);
|
||||
}
|
||||
|
||||
// Raw bytes of sub-block `ib` of the block at `bp` (IQ: the 256-value super-block, ib 0..7; Q2_0: the 64-value block,
|
||||
// ib 0..1; IQ4_NL: the 32-value block).
|
||||
template <int T> __device__ __forceinline__ void load_unit(const uint8_t* bp, int ib, uint32_t (&w)[5]) {
|
||||
if constexpr (T == T_IQ2_XXS) {
|
||||
w[0] = ld32(bp + 2 + 8 * ib); w[1] = ld32(bp + 6 + 8 * ib); w[2] = ld16(bp);
|
||||
} else if constexpr (T == T_IQ2_XS) {
|
||||
w[0] = ld32(bp + 2 + 8 * ib); w[1] = ld32(bp + 6 + 8 * ib); w[2] = ld16(bp) | ((uint32_t) bp[66 + ib] << 16);
|
||||
} else if constexpr (T == T_IQ2_S) {
|
||||
w[0] = ld32(bp + 2 + 4 * ib); w[1] = ld32(bp + 34 + 4 * ib);
|
||||
w[2] = ld16(bp) | ((uint32_t) bp[66 + ib] << 16) | ((uint32_t) bp[74 + ib] << 24);
|
||||
} else if constexpr (T == T_IQ3_XXS) {
|
||||
w[0] = ld32(bp + 2 + 8 * ib); w[1] = ld32(bp + 6 + 8 * ib); w[2] = ld32(bp + 66 + 4 * ib); w[3] = ld16(bp);
|
||||
} else if constexpr (T == T_IQ3_S) {
|
||||
w[0] = ld32(bp + 2 + 8 * ib); w[1] = ld32(bp + 6 + 8 * ib); w[2] = ld32(bp + 74 + 4 * ib);
|
||||
w[3] = ld16(bp) | ((uint32_t) bp[66 + ib] << 16) | ((uint32_t) ((bp[106 + ib / 2] >> (4 * (ib & 1))) & 15) << 24);
|
||||
} else if constexpr (T == T_IQ4_XS) {
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k) w[k] = ld32(bp + 8 + 16 * ib + 4 * k);
|
||||
const uint32_t ls = ((bp[4 + ib / 2] >> (4 * (ib & 1))) & 15) | (((ld16(bp + 2) >> (2 * ib)) & 3) << 4);
|
||||
w[4] = ld16(bp) | (ls << 16);
|
||||
} else if constexpr (T == T_IQ4_NL) {
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k) w[k] = ld32(bp + 2 + 4 * k);
|
||||
w[4] = ld16(bp);
|
||||
} else { // Q2_0
|
||||
w[0] = ld32(bp + 2 + 8 * ib); w[1] = ld32(bp + 6 + 8 * ib); w[2] = ld16(bp);
|
||||
}
|
||||
}
|
||||
|
||||
// The sub-block as 32 int8 (q[0..7], natural order) and its scales (s0: values 0-15, s1: 16-31).
|
||||
template <int T>
|
||||
__device__ __forceinline__ void convert(const uint32_t (&w)[5], const uint8_t* grid, const uint32_t (&kv)[4],
|
||||
uint32_t (&q)[8], float& s0, float& s1) {
|
||||
if constexpr (T == T_IQ2_XXS) {
|
||||
const uint2* g = (const uint2*) grid;
|
||||
#pragma unroll
|
||||
for (int l = 0; l < 4; ++l) {
|
||||
const uint2 e = g[(w[0] >> (8 * l)) & 255];
|
||||
signed8(e.x, e.y, unpack_ksigns(w[1] >> (7 * l)), q[2 * l], q[2 * l + 1]);
|
||||
}
|
||||
s0 = s1 = half_at(w[2]) * (float) ((w[1] >> 27) | 1) * 0.125f;
|
||||
} else if constexpr (T == T_IQ2_XS) {
|
||||
const uint2* g = (const uint2*) grid;
|
||||
#pragma unroll
|
||||
for (int l = 0; l < 4; ++l) {
|
||||
const uint32_t c = (w[l >> 1] >> (16 * (l & 1))) & 0xFFFF;
|
||||
const uint2 e = g[c & 511];
|
||||
signed8(e.x, e.y, unpack_ksigns(c >> 9), q[2 * l], q[2 * l + 1]);
|
||||
}
|
||||
const float d = half_at(w[2]);
|
||||
const uint32_t sc = w[2] >> 16;
|
||||
s0 = d * (float) (2 * (sc & 15) + 1) * 0.125f;
|
||||
s1 = d * (float) (2 * ((sc >> 4) & 15) + 1) * 0.125f;
|
||||
} else if constexpr (T == T_IQ2_S) {
|
||||
const uint2* g = (const uint2*) grid;
|
||||
const uint32_t qh = (w[2] >> 16) & 255;
|
||||
#pragma unroll
|
||||
for (int l = 0; l < 4; ++l) {
|
||||
const uint2 e = g[((w[0] >> (8 * l)) & 255) | ((qh << (8 - 2 * l)) & 0x300)];
|
||||
signed8(e.x, e.y, ((w[1] >> (8 * l)) & 255) * 0x01010101u, q[2 * l], q[2 * l + 1]);
|
||||
}
|
||||
const float d = half_at(w[2]);
|
||||
const uint32_t sc = w[2] >> 24;
|
||||
s0 = d * (float) (2 * (sc & 15) + 1) * 0.125f;
|
||||
s1 = d * (float) (2 * (sc >> 4) + 1) * 0.125f;
|
||||
} else if constexpr (T == T_IQ3_XXS) {
|
||||
const uint32_t* g = (const uint32_t*) grid;
|
||||
#pragma unroll
|
||||
for (int l = 0; l < 4; ++l) {
|
||||
const uint32_t i0 = (w[l >> 1] >> (16 * (l & 1))) & 255, i1 = (w[l >> 1] >> (16 * (l & 1) + 8)) & 255;
|
||||
signed8(g[i0], g[i1], unpack_ksigns(w[2] >> (7 * l)), q[2 * l], q[2 * l + 1]);
|
||||
}
|
||||
s0 = s1 = half_at(w[3]) * (float) (2 * (w[2] >> 28) + 1) * 0.25f;
|
||||
} else if constexpr (T == T_IQ3_S) {
|
||||
const uint32_t* g = (const uint32_t*) grid;
|
||||
const uint32_t qh = (w[3] >> 16) & 255;
|
||||
#pragma unroll
|
||||
for (int l = 0; l < 4; ++l) {
|
||||
const uint32_t i0 = (w[l >> 1] >> (16 * (l & 1))) & 255, i1 = (w[l >> 1] >> (16 * (l & 1) + 8)) & 255;
|
||||
signed8(g[i0 | ((qh << (8 - 2 * l)) & 256)], g[i1 | ((qh << (7 - 2 * l)) & 256)],
|
||||
((w[2] >> (8 * l)) & 255) * 0x01010101u, q[2 * l], q[2 * l + 1]);
|
||||
}
|
||||
s0 = s1 = half_at(w[3]) * (float) (1 + 2 * (w[3] >> 24));
|
||||
} else if constexpr (T == T_IQ4_XS || T == T_IQ4_NL) {
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k) table16(w[k], kv, q[k], q[4 + k]);
|
||||
s0 = s1 = T == T_IQ4_XS ? half_at(w[4]) * (float) ((int) (w[4] >> 16) - 32) : half_at(w[4]);
|
||||
} else { // Q2_0: code - 1 via a byte table {-1, 0, 1, 2}
|
||||
#pragma unroll
|
||||
for (int h = 0; h < 4; ++h) {
|
||||
const uint32_t c = (w[h >> 1] >> (16 * (h & 1))) & 0xFFFF;
|
||||
const uint32_t qe = __byte_perm(0x020100FFu, 0x020100FFu, c & 0x7777);
|
||||
const uint32_t qo = __byte_perm(0x020100FFu, 0x020100FFu, (c >> 2) & 0x7777);
|
||||
q[2 * h] = __byte_perm(qe, qo, 0x5140);
|
||||
q[2 * h + 1] = __byte_perm(qe, qo, 0x7362);
|
||||
}
|
||||
s0 = s1 = half_at(w[2]);
|
||||
}
|
||||
}
|
||||
|
||||
template <int T> __device__ __forceinline__ const void* grid_src() {
|
||||
if constexpr (T == T_IQ2_XXS) return iq2xxs_grid;
|
||||
else if constexpr (T == T_IQ2_XS) return iq2xs_grid;
|
||||
else if constexpr (T == T_IQ2_S) return iq2s_grid;
|
||||
else if constexpr (T == T_IQ3_XXS) return iq3xxs_grid;
|
||||
else if constexpr (T == T_IQ3_S) return iq3s_grid;
|
||||
else return nullptr;
|
||||
}
|
||||
#endif
|
||||
|
||||
// One work item = (TR routed rows of expert e - one or two of group()'s 64-row tiles -, a block of WR weight rows);
|
||||
// a persistent grid walks the batch's items. GU: gate/up rows of WR / 2 features (format WT) against the rows'
|
||||
// tokens' activations, SwiGLU, H to int8 per 32 features into `out`. !GU: down rows (format WT) against H, FP32 into
|
||||
// `dm`. Warp (wf, wt): weight rows 64 wf.. as four m16 tiles (GU: a tile = 8 features, gate rows as mma rows 0-7 and
|
||||
// up rows as 8-15, so a lane holds gate and up of one feature), routed rows 16 wt.. as two n8.
|
||||
template <int WT, bool GU, int WW>
|
||||
__global__ void __launch_bounds__(THREADS, 1)
|
||||
native_kernel(const Batch b, const NativeGeom geo, const Tables tb, const uint8_t* __restrict__ act,
|
||||
const int32_t* __restrict__ src, uint8_t* __restrict__ out, float* __restrict__ dm) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
constexpr int TR = tile_rows(WW), WR = weight_rows(WW);
|
||||
constexpr int WT_BYTES = WR * WLD, WS_FLOATS = WR * 4, ACT_STAGE = TR * AB;
|
||||
constexpr int NS = (GU ? GU_ROWS_K : D_ROWS_K) / 64; // 64-value stages along K
|
||||
constexpr int NFB = (GU ? 1280 : 2560) / WR; // weight-row blocks per tile
|
||||
constexpr int ACT_LD = NS * AB; // bytes per activation row: 3200 (a token), 800 (a row's H)
|
||||
constexpr int BS = block_bytes(WT);
|
||||
constexpr bool K16 = per16(WT);
|
||||
extern __shared__ __align__(16) uint8_t smem[];
|
||||
uint8_t* wt = smem; // [2][WR][WLD] decoded int8 weights
|
||||
float* ws = (float*) (smem + 2 * WT_BYTES); // [2][WR][4] their scales per 16 values
|
||||
uint8_t* stages = (uint8_t*) (ws + 2 * WS_FLOATS); // [ASTAGES][TR][AB] activations
|
||||
int* srow = (int*) (stages + ASTAGES * ACT_STAGE); // [TR] the activation row of each tile row
|
||||
uint8_t* sgrid = (uint8_t*) (srow + TR); // the codebook
|
||||
|
||||
const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, tig = lane & 3;
|
||||
const int wf = warp % WW, nb0 = 16 * (warp / WW);
|
||||
if constexpr (grid_bytes(WT) > 0) {
|
||||
const uint32_t* gs = (const uint32_t*) grid_src<WT>();
|
||||
for (int i = tid; i < grid_bytes(WT) / 4; i += THREADS) ((uint32_t*) sgrid)[i] = gs[i];
|
||||
}
|
||||
uint32_t kv[4] = {0, 0, 0, 0};
|
||||
if constexpr (WT == T_IQ4_XS || WT == T_IQ4_NL) {
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 16; ++k) kv[k >> 2] |= (uint32_t) (uint8_t) kvalues_iq4nl[k] << (8 * (k & 3));
|
||||
}
|
||||
// Local weight rows: m tile i of warp wf is rows 64 wf + 16 i.. (mma rows g and g + 8: rows +g and +8+g). For
|
||||
// gate/up a tile is 8 features, its gate rows then its up rows, so a lane holds gate and up of one feature. The
|
||||
// fragments come by ldmatrix: lane l gives the address of row (l & 7) of matrix l >> 3 - A: {rows 0-7, rows 8-15}
|
||||
// x {bytes 0-15, 16-31} of a 32-value chunk; B: {tokens nb0 + 0-7, + 8-15} x {bytes 0-15, 16-31}.
|
||||
const int lm = lane >> 3, l8 = lane & 7;
|
||||
const int a_off = (64 * wf + 8 * (lm & 1) + l8) * WLD + 16 * (lm >> 1);
|
||||
const int b_off = (nb0 + 8 * (lm >> 1) + l8) * AB + 16 * (lm & 1);
|
||||
// this thread's decode unit (the first 2 WR threads): local weight row ur, sub-block uj of each stage
|
||||
const bool dec = tid < 2 * WR;
|
||||
const int ur = dec ? tid >> 1 : 0, uj = tid & 1;
|
||||
const int t0 = tb.ts[b.e0], nwork = (tb.ts[b.e1] - t0) * NFB;
|
||||
for (int w = blockIdx.x; w < nwork; w += gridDim.x) {
|
||||
const int2 tl = tb.tiles[t0 + w / NFB];
|
||||
const int fb = w % NFB, e = tl.x, row0 = tl.y;
|
||||
if ((row0 - tb.off[e]) % TR != 0) continue; // a 64-row tile inside an earlier item's TR rows
|
||||
const int nrows = min(TR, tb.off[e + 1] - row0);
|
||||
const uint8_t* blob = b.blob[e - b.e0];
|
||||
const int rbase = fb * WR;
|
||||
// local row ur: gate/up - feature (WR / 2) fb + 8 (ur >> 4) + (ur & 7), its up row when bit 3 is set
|
||||
const uint8_t* wrow = GU ? blob + ((ur & 8) ? geo.up_off : 0) +
|
||||
(size_t) (fb * (WR / 2) + 8 * (ur >> 4) + (ur & 7)) * geo.gu_row
|
||||
: blob + geo.down_off + (size_t) (rbase + ur) * geo.d_row;
|
||||
auto unit = [&](int s) -> const uint8_t* { // the block of stage s's sub-block
|
||||
if (GU) return wrow + (s >> 2) * BS;
|
||||
return WT == T_IQ4_NL ? wrow + (2 * s + uj) * BS : wrow + s * BS;
|
||||
};
|
||||
auto sub = [&](int s) { return GU ? 2 * (s & 3) + uj : uj; };
|
||||
// the weight bytes into L2 ahead of the register loads (one stage ahead only - less than a DRAM round trip):
|
||||
// gate/up PF super-blocks ahead (the row's two threads take a block's first and last line), down the whole row
|
||||
// slice at once (180 or 360 bytes a row)
|
||||
constexpr int PF = 2;
|
||||
auto prefetch_sb = [&](int sb) {
|
||||
if (dec && sb < GU_ROWS_K / 256) pf_l2(wrow + sb * BS + (uj ? BS - 1 : 0));
|
||||
};
|
||||
if (GU) {
|
||||
#pragma unroll
|
||||
for (int sb = 0; sb < PF; ++sb) prefetch_sb(sb);
|
||||
} else if (dec) {
|
||||
for (size_t o = 128 * (size_t) uj; o < geo.d_row + 127; o += 256) pf_l2(wrow + min(o, geo.d_row - 1));
|
||||
}
|
||||
__syncthreads(); // the previous item is done with the buffers
|
||||
if (tid < TR) srow[tid] = tid < nrows ? (GU ? src[row0 + tid] : row0 + tid) : -1;
|
||||
__syncthreads(); // srow
|
||||
// the rows past the item's end are not loaded: a warp entirely past it skips the products, and the columns
|
||||
// of a partial one are never written (no reduction mixes columns)
|
||||
auto load_act = [&](int s) {
|
||||
uint8_t* st = stages + (s % ASTAGES) * ACT_STAGE;
|
||||
for (int c = tid; c < TR * 5; c += THREADS) {
|
||||
const int r = c / 5, q = c % 5;
|
||||
if (r < nrows) cp16(st + r * AB + q * 16, act + (size_t) srow[r] * ACT_LD + s * AB + q * 16);
|
||||
}
|
||||
};
|
||||
auto put = [&](const uint32_t (&raw)[5], int buf) {
|
||||
if (!dec) return;
|
||||
uint32_t q[8];
|
||||
float s0, s1;
|
||||
convert<WT>(raw, sgrid, kv, q, s0, s1);
|
||||
uint4* d = (uint4*) (wt + buf * WT_BYTES + ur * WLD + 32 * uj);
|
||||
d[0] = make_uint4(q[0], q[1], q[2], q[3]);
|
||||
d[1] = make_uint4(q[4], q[5], q[6], q[7]);
|
||||
*(float2*) (ws + buf * WS_FLOATS + ur * 4 + 2 * uj) = make_float2(s0, s1);
|
||||
};
|
||||
#pragma unroll
|
||||
for (int s = 0; s < ASTAGES - 1; ++s) {
|
||||
if (s < NS) load_act(s);
|
||||
cp_commit();
|
||||
}
|
||||
uint32_t raw[5] = {0, 0, 0, 0, 0};
|
||||
if (dec) load_unit<WT>(unit(0), sub(0), raw);
|
||||
put(raw, 0);
|
||||
if (dec) load_unit<WT>(unit(1), sub(1), raw);
|
||||
// a warp whose rows are all past the tile's end only takes part in the loads
|
||||
const bool on0 = nb0 < nrows, on1 = nb0 + 8 < nrows;
|
||||
float acc[4][2][4];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i)
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 2; ++n)
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 4; ++q) acc[i][n][q] = 0.0f;
|
||||
for (int s = 0; s < NS; ++s) {
|
||||
cp_wait<ASTAGES - 2>();
|
||||
__syncthreads();
|
||||
if (s + ASTAGES - 1 < NS) load_act(s + ASTAGES - 1);
|
||||
cp_commit();
|
||||
if (GU && (s & 3) == 0) prefetch_sb((s >> 2) + PF);
|
||||
if (on0) {
|
||||
const uint8_t* W = wt + (s & 1) * WT_BYTES;
|
||||
const float* S = ws + (s & 1) * WS_FLOATS;
|
||||
const uint8_t* sa = stages + (s % ASTAGES) * ACT_STAGE;
|
||||
// the stage's scales: per weight row 4 (per 16 values), per routed row 2 (per 32)
|
||||
float2 dx[2][2]; // [n8 tile][column 2 tig + cc]
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 2; ++n)
|
||||
#pragma unroll
|
||||
for (int cc = 0; cc < 2; ++cc) dx[n][cc] = *(const float2*) (sa + (nb0 + 8 * n + 2 * tig + cc) * AB + 64);
|
||||
#pragma unroll
|
||||
for (int h = 0; h < 2; ++h) { // the stage's two 32-value halves
|
||||
uint32_t bq[4]; // {n0 bytes 0-15, n0 16-31, n1 0-15, n1 16-31}
|
||||
ldsm4(bq, sa + b_off + 32 * h);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
uint32_t a[4]; // {rows g, g + 8} x {bytes 0-15, 16-31}
|
||||
ldsm4(a, W + a_off + 16 * i * WLD + 32 * h);
|
||||
const float2 swa = *(const float2*) (S + (64 * wf + 16 * i + g) * 4 + 2 * h);
|
||||
const float2 swb = *(const float2*) (S + (64 * wf + 16 * i + 8 + g) * 4 + 2 * h);
|
||||
const float wa0 = swa.x, wa1 = swa.y, wb0 = swb.x, wb1 = swb.y;
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 2; ++n) {
|
||||
if (n == 1 && !on1) break;
|
||||
if constexpr (K16) {
|
||||
// a scale per 16 weights: the two halves' dots, each times its scale, then d_x
|
||||
int d0[4], d1[4];
|
||||
mma16(d0, a[0], a[1], bq[2 * n]);
|
||||
mma16(d1, a[2], a[3], bq[2 * n + 1]);
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 4; ++q) {
|
||||
const float v = fmaf(q < 2 ? wa1 : wb1, dotf(d1[q]), (q < 2 ? wa0 : wb0) * dotf(d0[q]));
|
||||
acc[i][n][q] = fmaf(h ? dx[n][q & 1].y : dx[n][q & 1].x, v, acc[i][n][q]);
|
||||
}
|
||||
} else {
|
||||
int d[4];
|
||||
mma32(d, a, bq[2 * n], bq[2 * n + 1]);
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 4; ++q) {
|
||||
const float p = (q < 2 ? wa0 : wb0) * (h ? dx[n][q & 1].y : dx[n][q & 1].x);
|
||||
acc[i][n][q] = fmaf(p, dotf(d[q]), acc[i][n][q]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// the next stage's weights (into the other buffer, read after the next barrier). (Decoding them before
|
||||
// this stage's products instead was slower on the RTX 5070: a layer of 2048 / 3584 / 8192 tokens, IQ2_S
|
||||
// 5.2 / 9.5 / 17.5 ms -> 6.0 / 10.2 / 18.3.)
|
||||
if (s + 1 < NS) {
|
||||
put(raw, (s + 1) & 1);
|
||||
if (dec && s + 2 < NS) load_unit<WT>(unit(s + 2), sub(s + 2), raw);
|
||||
}
|
||||
}
|
||||
if (!on0) continue;
|
||||
if (GU) {
|
||||
// SwiGLU of feature (WR / 2) fb + 32 wf + 8 i + g (lane g of m tile i: rows 64 wf + 16 i + g and + 8), tile
|
||||
// rows nb0 + 8 n + 2 tig + c; H to int8 per row over the warp's 32 features (one 32-value half-block of the
|
||||
// down product's K)
|
||||
const int blk = WW * fb + wf;
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 2; ++n) {
|
||||
if (n == 1 && !on1) break;
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 2; ++c) {
|
||||
float h[4], am = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
const float gt = acc[i][n][c], up = acc[i][n][2 + c];
|
||||
h[i] = gt / (1.0f + __expf(-gt)) * up;
|
||||
am = fmaxf(am, fabsf(h[i]));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int o = 4; o < 32; o <<= 1) am = fmaxf(am, __shfl_xor_sync(0xffffffffu, am, o));
|
||||
const float inv = am > 0.0f ? 127.0f / am : 0.0f;
|
||||
const int r = nb0 + 8 * n + 2 * tig + c;
|
||||
if (r < nrows) {
|
||||
uint8_t* o = out + (size_t) (row0 + r) * (10 * AB) + (blk >> 1) * AB;
|
||||
const int hh = blk & 1;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) o[32 * hh + 8 * i + g] = (uint8_t) (int8_t) __float2int_rn(h[i] * inv);
|
||||
if (g == 0) *(float*) (o + 64 + 4 * hh) = am / 127.0f;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i)
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 2; ++n) {
|
||||
if (n == 1 && !on1) break;
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 4; ++q) {
|
||||
const int r = nb0 + 8 * n + 2 * tig + (q & 1);
|
||||
if (r < nrows) dm[(size_t) (row0 + r) * 2560 + rbase + 64 * wf + 16 * i + 8 * (q >> 1) + g] = acc[i][n][q];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
unsigned blocks(int64_t n, int per) { return (unsigned) ((n + per - 1) / per); }
|
||||
|
||||
// per device: whether every kernel here runs (sm_80+, device code in this build, fits), and their occupancy
|
||||
struct DevInfo {
|
||||
bool done = false, ok = false;
|
||||
int sms = 0, occ = 1;
|
||||
};
|
||||
std::mutex g_mu;
|
||||
DevInfo g_dev[32];
|
||||
|
||||
template <int T, bool GU, int WW> bool setup_ww(int& occ) {
|
||||
cudaFuncAttributes fa{};
|
||||
if (cudaFuncGetAttributes(&fa, native_kernel<T, GU, WW>) != cudaSuccess || fa.ptxVersion < 80) return false;
|
||||
if (cudaFuncSetAttribute(native_kernel<T, GU, WW>, cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
(int) smem_bytes(T, WW)) != cudaSuccess)
|
||||
return false;
|
||||
int o = 0;
|
||||
if (cudaOccupancyMaxActiveBlocksPerMultiprocessor(&o, native_kernel<T, GU, WW>, THREADS, smem_bytes(T, WW)) !=
|
||||
cudaSuccess || o < 1)
|
||||
return false;
|
||||
occ = std::min(occ, o);
|
||||
return true;
|
||||
}
|
||||
template <int T, bool GU> bool setup_one(int& occ) { return setup_ww<T, GU, 4>(occ) && setup_ww<T, GU, 2>(occ); }
|
||||
|
||||
const DevInfo& dev_info() {
|
||||
int dev = 0;
|
||||
cudaGetDevice(&dev);
|
||||
std::lock_guard<std::mutex> lk(g_mu);
|
||||
DevInfo& d = g_dev[dev & 31];
|
||||
if (d.done) return d;
|
||||
d.done = true;
|
||||
int major = 0;
|
||||
cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev);
|
||||
cudaDeviceGetAttribute(&d.sms, cudaDevAttrMultiProcessorCount, dev);
|
||||
if (major < 8) return d;
|
||||
int occ = 1 << 20;
|
||||
d.ok = setup_one<T_IQ2_XXS, true>(occ) && setup_one<T_IQ2_XS, true>(occ) && setup_one<T_IQ2_S, true>(occ) &&
|
||||
setup_one<T_IQ3_XXS, true>(occ) && setup_one<T_IQ3_S, true>(occ) && setup_one<T_IQ4_XS, true>(occ) &&
|
||||
setup_one<T_Q2_0, false>(occ) && setup_one<T_IQ4_NL, false>(occ);
|
||||
d.occ = d.ok ? occ : 1;
|
||||
cudaGetLastError();
|
||||
return d;
|
||||
}
|
||||
|
||||
bool gu_covered(int t) {
|
||||
return t == T_IQ2_XXS || t == T_IQ2_XS || t == T_IQ2_S || t == T_IQ3_XXS || t == T_IQ3_S || t == T_IQ4_XS;
|
||||
}
|
||||
bool d_covered(int t) { return t == T_Q2_0 || t == T_IQ4_NL; }
|
||||
|
||||
template <int T, bool GU>
|
||||
void launch(int ww, unsigned grid, const Batch& b, const NativeGeom& g, const Tables& tb, const void* act,
|
||||
const int32_t* src, void* out, float* dm, cudaStream_t s) {
|
||||
const uint8_t* a = (const uint8_t*) act;
|
||||
uint8_t* o = (uint8_t*) out;
|
||||
if (ww == 4) native_kernel<T, GU, 4><<<grid, THREADS, smem_bytes(T, 4), s>>>(b, g, tb, a, src, o, dm);
|
||||
else native_kernel<T, GU, 2><<<grid, THREADS, smem_bytes(T, 2), s>>>(b, g, tb, a, src, o, dm);
|
||||
}
|
||||
|
||||
// The work item's shape for a layer of `n` routed rows over `n_expert` experts: 128 routed rows when an expert has
|
||||
// about one such tile (56 to 112 rows on average), else 64. Measured on the RTX 5070 (tests/cuda/prefill_fused_iq_test,
|
||||
// a layer of 2048 / 3584 / 8192 tokens, ms, 64 -> 128): IQ2_S 5.2 -> 6.7, 9.5 -> 9.3, 17.5 -> 19.5; IQ2_XXS 4.2 ->
|
||||
// 5.8, 7.4 -> 7.6, 13.6 -> 16.5; IQ3_S 5.6 -> 6.6, 10.2 -> 8.8, 18.4 -> 19.0; IQ3_XXS 5.1 -> 6.3, 9.1 -> 8.3, 16.5 ->
|
||||
// 17.5 (256 rows was slower still). STRATA_PF_FUSED_TILE=64|128 forces one.
|
||||
int pick_ww(int64_t n, int n_expert) {
|
||||
static const int forced = [] {
|
||||
const char* v = std::getenv("STRATA_PF_FUSED_TILE");
|
||||
const int t = v ? std::atoi(v) : 0;
|
||||
return t == 64 ? 4 : t == 128 ? 2 : 0;
|
||||
}();
|
||||
if (forced) return forced;
|
||||
const double avg = (double) n / std::max(n_expert, 1);
|
||||
return avg > 56.0 && avg <= 112.0 ? 2 : 4;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
bool native_supported(int gu_type, int d_type) {
|
||||
static const bool off = [] { // STRATA_PF_FUSED_NATIVE=0: the native packs keep MMQ under STRATA_PF_FUSED=1 (A/B)
|
||||
const char* v = std::getenv("STRATA_PF_FUSED_NATIVE");
|
||||
return v != nullptr && v[0] == '0';
|
||||
}();
|
||||
return !off && requested() && gu_covered(gu_type) && d_covered(d_type) && dev_info().ok; // opt-in (=1)
|
||||
}
|
||||
|
||||
void quantize_act_native(const float* x, int64_t rows, int64_t cols, void* xa, void* stream) {
|
||||
if (rows <= 0) return;
|
||||
const int64_t nblk = rows * (cols / 64);
|
||||
quant_act_nat_kernel<<<blocks(nblk, 8), 256, 0, (cudaStream_t) stream>>>(x, nblk, (uint8_t*) xa);
|
||||
ck(cudaGetLastError(), "quantize_act_native");
|
||||
}
|
||||
|
||||
void experts_native(const Batch& b, const NativeGeom& g, int n_expert, int64_t n, const void* scratch, const void* xa,
|
||||
const int32_t* src, void* ha, float* dm, void* stream) {
|
||||
if (b.e1 <= b.e0 || n <= 0) return;
|
||||
const DevInfo& d = dev_info();
|
||||
if (!d.ok || !gu_covered(g.gu_type) || !d_covered(g.d_type)) {
|
||||
std::fprintf(stderr, "prefill fused experts (native): types %d / %d are not covered here\n", g.gu_type, g.d_type);
|
||||
std::exit(1);
|
||||
}
|
||||
const cudaStream_t s = (cudaStream_t) stream;
|
||||
const Tables tb = tables(const_cast<void*>(scratch), n_expert);
|
||||
const int ww = pick_ww(n, n_expert);
|
||||
// the most 64-row tiles the batch can have (every row in it, plus a partial tile per expert) - the items of the
|
||||
// ones inside a larger item end at once
|
||||
const int64_t tiles = (n + kTileRows - 1) / kTileRows + (b.e1 - b.e0);
|
||||
const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (1280 / weight_rows(ww)), (int64_t) d.sms * d.occ);
|
||||
const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (2560 / weight_rows(ww)), (int64_t) d.sms * d.occ);
|
||||
switch (g.gu_type) {
|
||||
case T_IQ2_XXS: launch<T_IQ2_XXS, true>(ww, g_gu, b, g, tb, xa, src, ha, nullptr, s); break;
|
||||
case T_IQ2_XS: launch<T_IQ2_XS, true>(ww, g_gu, b, g, tb, xa, src, ha, nullptr, s); break;
|
||||
case T_IQ2_S: launch<T_IQ2_S, true>(ww, g_gu, b, g, tb, xa, src, ha, nullptr, s); break;
|
||||
case T_IQ3_XXS: launch<T_IQ3_XXS, true>(ww, g_gu, b, g, tb, xa, src, ha, nullptr, s); break;
|
||||
case T_IQ3_S: launch<T_IQ3_S, true>(ww, g_gu, b, g, tb, xa, src, ha, nullptr, s); break;
|
||||
default: launch<T_IQ4_XS, true>(ww, g_gu, b, g, tb, xa, src, ha, nullptr, s); break;
|
||||
}
|
||||
if (g.d_type == T_Q2_0) launch<T_Q2_0, false>(ww, g_d, b, g, tb, ha, src, nullptr, dm, s);
|
||||
else launch<T_IQ4_NL, false>(ww, g_d, b, g, tb, ha, src, nullptr, dm, s);
|
||||
ck(cudaGetLastError(), "experts_native");
|
||||
}
|
||||
|
||||
} // namespace strata::prefill::fused
|
||||
+372
-236
@@ -23,6 +23,8 @@
|
||||
#include "strata/kernels/qsa_prompt_attn.hpp"
|
||||
#include "strata/kernels/qsa_select.hpp"
|
||||
#include "strata/prefill/gemm.hpp"
|
||||
#include "strata/prefill/moe_fused.hpp"
|
||||
#include "strata/prefill/moe_fused_iq.hpp"
|
||||
#include "strata/prefill/moe_mmq.hpp"
|
||||
#include "strata/prefill/kernels.hpp"
|
||||
|
||||
@@ -59,6 +61,24 @@ void swiglu(const float*, float*, int64_t, int64_t, bool, void*) {}
|
||||
void iota(int32_t*, int64_t, void*) {}
|
||||
} // namespace strata::prefill::mmq
|
||||
#endif
|
||||
#ifndef STRATA_PREFILL_FUSED
|
||||
// #136: the fused int8 experts are CUDA-only (HIP and builds without MMQ keep the MMQ / FP16 paths)
|
||||
namespace strata::prefill::fused {
|
||||
bool built() { return false; }
|
||||
bool available() { return false; }
|
||||
bool enabled() { return false; }
|
||||
bool requested() { return false; }
|
||||
size_t act_bytes(int64_t, int64_t) { return 0; }
|
||||
size_t group_bytes(int64_t, int) { return 0; }
|
||||
void quantize_act(const float*, int64_t, int64_t, void*, void*) {}
|
||||
void group(const int32_t*, int64_t, int, int, void*, int32_t*, int32_t*, void*) {}
|
||||
void experts(const Batch&, int, int64_t, const void*, const void*, const int32_t*, void*, float*, void*) {}
|
||||
bool native_supported(int, int) { return false; }
|
||||
void quantize_act_native(const float*, int64_t, int64_t, void*, void*) {}
|
||||
void experts_native(const Batch&, const NativeGeom&, int, int64_t, const void*, const void*, const int32_t*, void*,
|
||||
float*, void*) {}
|
||||
} // namespace strata::prefill::fused
|
||||
#endif
|
||||
|
||||
namespace strata::prefill {
|
||||
namespace {
|
||||
@@ -71,9 +91,9 @@ constexpr int64_t C = 10240, ZV = 6144, HV = 48;
|
||||
inline int64_t MAXBLOB() { return (int64_t) strata::kernels::cpu::expert_layout().max_blob; }
|
||||
constexpr int STAGE = 8; // host->device expert staging ring (chunks below stream_all_min())
|
||||
// Step 3: from this chunk size on, every non-resident expert of every layer streams in a fixed order through a
|
||||
// RING_MAX-slot ring (nearly all 512 are routed at such a chunk), so the copy engine keeps working through the
|
||||
// ring_slots()-slot ring (nearly all 512 are routed at such a chunk), so the copy engine keeps working through the
|
||||
// attention halves instead of waiting for each layer's routing.
|
||||
constexpr int RING_MAX = 512; // the arrays; the ring itself is ring_slots()
|
||||
constexpr int RING_MAX = 1024; // the arrays; the ring itself is ring_slots(), at most ring_cap()
|
||||
// The chunk size from which every expert streams: 1024 since 0.1.30 (was 2048). Measured on the 5070, Q2_0 / IQ2_XS,
|
||||
// fixed cache: 1,500-token prompts 621 -> 785 / 612 -> 735 tok/s, 2,000 727 -> 934 / 712 -> 892, 4,000 (its last
|
||||
// chunk) 779 -> 912 / 766 -> 844, the same output. Below ~1,000 tokens the output changed on Q2_0 (a smaller chunk
|
||||
@@ -88,6 +108,28 @@ double g_pinned_share = 1.0;
|
||||
// and 96 when a large share goes through host copies (IQ3_S on 64 GB, a third unpinned: 96 slots 1216, 256 1070 -
|
||||
// the host copies are the limit and the bigger ring only takes cache slots). STRATA_PREFILL_RING overrides.
|
||||
int g_ring_override = 0; // #340: set by a layer split (Prefill::set_ring_override); 0 = the rule below
|
||||
// #136: the fused experts (STRATA_PF_FUSED=1) launch on a batch of a layer's streamed experts at once, so the ring
|
||||
// should hold a whole layer's (~460 of 512 on Q2_0): with 384 slots a layer's last batch waits for slots its own
|
||||
// first batch frees. Measured on the 5070, Q2_0, the 4K / 32K code-agent prompts (one run each): fused at 384 slots
|
||||
// +5% / +3% over MMQ, at 512 +19-22% / +11-12% (MMQ itself at 512: -2% / -1%). P3: with the fused path's smaller
|
||||
// buffers (moe_bufs) 1024 slots - two layers' experts - fit too; 2 pairs each, prompt tok/s at 512 / 1024 slots: 4K
|
||||
// 1,512 / 1,605, 32K 2,485 / 2,660 (both chunk 8192), 128K with KV streaming 2,346 / 2,392 (the chunk falls from
|
||||
// 8192 to 6144, but more experts stay resident). A native pack likewise when the native kernels (moe_fused_iq.hpp)
|
||||
// take any of its layers: IQ2_XS, 4K / 32K, their first version at 384 slots -3% / -6% against MMQ, at 512 +8% / 0%.
|
||||
inline bool fused_ring() {
|
||||
if (!fused::enabled()) return false;
|
||||
const strata::kernels::cpu::ExpertLayout& lay = strata::kernels::cpu::expert_layout();
|
||||
if (!lay.native) return true;
|
||||
static const bool any = [&lay] {
|
||||
for (const auto& f : lay.fmt)
|
||||
if (fused::native_supported(f.gu_type, f.d_type)) return true;
|
||||
return false;
|
||||
}();
|
||||
return any;
|
||||
}
|
||||
// the largest ring: 512 slots; 1024 with the Q2_0 pack's fused experts (P3's smaller buffers, measured there) - the
|
||||
// native packs' fused layers were measured at 512
|
||||
inline int ring_cap() { return fused_ring() && !strata::kernels::cpu::expert_layout().native ? RING_MAX : 512; }
|
||||
inline int ring_slots(size_t T) {
|
||||
const char* v = std::getenv("STRATA_PREFILL_RING");
|
||||
#if defined(STRATA_USE_HIP)
|
||||
@@ -99,9 +141,10 @@ inline int ring_slots(size_t T) {
|
||||
}();
|
||||
if (!v && g_ring_override <= 0 && wmma) return (int64_t) T >= stream_all_min() ? 96 : STAGE;
|
||||
#endif
|
||||
const int r = v ? std::atoi(v) : g_ring_override > 0 ? g_ring_override : (g_pinned_share >= 0.9 ? 384 : 96);
|
||||
const int pinned_ring = fused_ring() ? 1024 : 384;
|
||||
const int r = v ? std::atoi(v) : g_ring_override > 0 ? g_ring_override : (g_pinned_share >= 0.9 ? pinned_ring : 96);
|
||||
if (v && r == STAGE) return STAGE; // Explicit opt-in to routed-only staging, including large chunks.
|
||||
const int big = r < 16 ? 16 : r > RING_MAX ? RING_MAX : r;
|
||||
const int big = r < 16 ? 16 : r > ring_cap() ? ring_cap() : r;
|
||||
return (int64_t) T >= stream_all_min() ? big : STAGE;
|
||||
}
|
||||
constexpr int DQ = 2; // dequantized-expert ring (FP16 gate/up + down)
|
||||
@@ -491,20 +534,41 @@ const MmqPlan& mmq_plan() {
|
||||
}();
|
||||
return plan;
|
||||
}
|
||||
uint64_t moe_set_bytes(size_t T, int64_t n_expert) {
|
||||
// #136 P3: a layout whose chunks run the fused experts (STRATA_PF_FUSED=1, the Q2_0 pack, a streamed chunk of
|
||||
// stream_all_min() tokens or more). Its GU, H and Xq hold only the fused path's grouping tables, int8 H and per-token
|
||||
// int8 activations, and Hq nothing: ~100 KB a token less than MMQ's FP32 GU / H and per-slot q8_1 rows, which is what
|
||||
// lets a bigger chunk or ring fit in the slots the prompt path borrows. A last chunk below stream_all_min() still runs
|
||||
// MMQ in the same buffers, so each keeps MMQ's size for stream_all_min() - 1 tokens. `src`: the layout streams experts
|
||||
// (Prefill::init got an ExpertSource; without one no chunk takes the streamed walk, so no chunk is fused).
|
||||
bool fused_layout(size_t T, bool src) {
|
||||
return src && fused_ring() && mmq_plan().any && ring_slots(T) > STAGE && (int64_t) T >= stream_all_min();
|
||||
}
|
||||
// The MoE buffers MMQ and the fused path share: GU and H in floats, Xq and Hq in bytes. Without `fused` (the
|
||||
// default): MMQ's, for T tokens.
|
||||
struct MoeBufs { size_t gu, h, xq, hq; };
|
||||
MoeBufs moe_bufs(size_t T, int64_t n_expert, bool fused) {
|
||||
if (!fused) return {T * K * 1280, T * K * 640, mmq::q8_bytes((int64_t) (T * K), N), mmq::q8_bytes((int64_t) (T * K), 640)};
|
||||
const size_t ts = (size_t) std::min<int64_t>((int64_t) T, stream_all_min() - 1); // MMQ's last small chunk
|
||||
return {std::max(ts * K * 1280, (fused::group_bytes((int64_t) (T * K), (int) n_expert) + 3) / 4),
|
||||
std::max(ts * K * 640, (fused::act_bytes((int64_t) (T * K), 640) + 3) / 4),
|
||||
std::max(mmq::q8_bytes((int64_t) (ts * K), N), fused::act_bytes((int64_t) T, N)),
|
||||
mmq::q8_bytes((int64_t) (ts * K), 640)};
|
||||
}
|
||||
uint64_t moe_set_bytes(size_t T, int64_t n_expert, bool fused) {
|
||||
const MmqPlan& mp = mmq_plan();
|
||||
const MoeBufs mb = moe_bufs(T, n_expert, fused);
|
||||
Alloc a; a.count_only = true; bool ok = true;
|
||||
a.take<float>(T * n_expert, ok); a.take<float>(T * K, ok); a.take<int32_t>(T * K, ok); a.take<int32_t>(T * K, ok);
|
||||
a.take<int32_t>(T * K, ok);
|
||||
if (mp.fallback) a.take<uint16_t>(T * K * N, ok);
|
||||
a.take<float>(T * K * 1280, ok);
|
||||
a.take<float>(mb.gu, ok);
|
||||
if (mp.fallback) a.take<uint16_t>(T * K * 640, ok);
|
||||
a.take<float>(T * K * N, ok); a.take<float>(T * 640, ok);
|
||||
a.take<float>(T * 640, ok); a.take<uint16_t>(T * 640, ok); a.take<float>(T * N, ok); a.take<float>(T, ok);
|
||||
if (mp.any) {
|
||||
a.take<uint8_t>(mmq::q8_bytes((int64_t) (T * K), N), ok);
|
||||
a.take<float>(T * K * 640, ok);
|
||||
a.take<uint8_t>(mmq::q8_bytes((int64_t) (T * K), 640), ok);
|
||||
a.take<uint8_t>(mb.xq, ok);
|
||||
a.take<float>(mb.h, ok);
|
||||
a.take<uint8_t>(mb.hq, ok);
|
||||
}
|
||||
return a.used;
|
||||
}
|
||||
@@ -544,7 +608,7 @@ bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, co
|
||||
m.borrowed = borrow != nullptr;
|
||||
bool ok = true;
|
||||
// one-time: events, the stager, the host buffers (for the largest chunk), the identity page table
|
||||
for (int i = 0; i < RING_MAX; ++i) {
|
||||
for (int i = 0; i < ring_cap(); ++i) {
|
||||
if (cudaEventCreateWithFlags(&m.copied[i], cudaEventDisableTiming) != cudaSuccess) ok = false;
|
||||
if (cudaEventCreateWithFlags(&m.used[i], cudaEventDisableTiming) != cudaSuccess) ok = false;
|
||||
}
|
||||
@@ -648,8 +712,10 @@ bool Prefill::carve(size_t T, void* alloc) {
|
||||
m.max_blocks = ss.qsa_states[ss.qsa_primary()].max_cells / s.idx_block + 2;
|
||||
{
|
||||
// one region for the attention half's and the MoE half's scratch (see gdn_set_bytes)
|
||||
const bool fz = fused_layout(T, m.src != nullptr);
|
||||
const MoeBufs mb = moe_bufs(T, m.g->n_expert, fz);
|
||||
const uint64_t region = std::max({gdn_set_bytes(T), qsa_set_bytes(T, m.cap, m.max_blocks, m.sel_batch,
|
||||
m.attn_batch, s), moe_set_bytes(T, m.g->n_expert)});
|
||||
m.attn_batch, s), moe_set_bytes(T, m.g->n_expert, fz)});
|
||||
uint8_t* base = o.take<uint8_t>((size_t) region, ok);
|
||||
m.region = base;
|
||||
m.region_bytes = region;
|
||||
@@ -672,15 +738,15 @@ bool Prefill::carve(size_t T, void* alloc) {
|
||||
m.slot_dev = c.take<int32_t>(T * K, ok); m.src_dev = c.take<int32_t>(T * K, ok);
|
||||
const MmqPlan& mp = mmq_plan();
|
||||
m.Xs = mp.fallback ? c.take<uint16_t>(T * K * N, ok) : nullptr;
|
||||
m.GU = c.take<float>(T * K * 1280, ok);
|
||||
m.GU = c.take<float>(mb.gu, ok);
|
||||
m.Hh = mp.fallback ? c.take<uint16_t>(T * K * 640, ok) : nullptr;
|
||||
m.Dm = c.take<float>(T * K * N, ok);
|
||||
m.sgate = c.take<float>(T * 640, ok); m.sup = c.take<float>(T * 640, ok); m.sh_h = c.take<uint16_t>(T * 640, ok);
|
||||
m.shared = c.take<float>(T * N, ok); m.sg = c.take<float>(T, ok);
|
||||
if (mp.any) {
|
||||
m.Xq = c.take<uint8_t>(mmq::q8_bytes((int64_t) (T * K), N), ok);
|
||||
m.H = c.take<float>(T * K * 640, ok);
|
||||
m.Hq = c.take<uint8_t>(mmq::q8_bytes((int64_t) (T * K), 640), ok);
|
||||
m.Xq = c.take<uint8_t>(mb.xq, ok);
|
||||
m.H = c.take<float>(mb.h, ok);
|
||||
m.Hq = c.take<uint8_t>(mb.hq, ok);
|
||||
}
|
||||
if (base == nullptr) ok = false;
|
||||
}
|
||||
@@ -904,7 +970,7 @@ uint64_t Prefill::bytes_needed(const core::ModelGeometry& g, const core::Session
|
||||
const int64_t cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s);
|
||||
const int64_t max_blocks = ss.qsa_states[ss.qsa_primary()].max_cells / s.idx_block + 2;
|
||||
o.take<uint8_t>((size_t) std::max({gdn_set_bytes(T), qsa_set_bytes(T, cap, max_blocks, 256, 32, s),
|
||||
moe_set_bytes(T, g.n_expert)}), ok);
|
||||
moe_set_bytes(T, g.n_expert, fused_layout(T, true))}), ok);
|
||||
for (int i = 0; i < DQ; ++i) { o.take<uint16_t>(1280 * 2560, ok); o.take<uint16_t>(2560 * 640, ok); }
|
||||
if (mmq_plan().any) {
|
||||
const MmqPlan& mp = mmq_plan();
|
||||
@@ -1579,235 +1645,303 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
|
||||
if (wgi->kind != core::WeightKind::Bf16InF32) { err = "prefill: shared gate is not BF16"; return false; }
|
||||
m.gemm.bf16(m.mixed_bf, (const uint16_t*) wgi->data, m.sg, T, 1, N);
|
||||
if (m.mixed_bf_lo) m.gemm.bf16(m.mixed_bf_lo, (const uint16_t*) wgi->data, m.sg, T, 1, N, 0, 1.0f);
|
||||
// group the (token, k) pairs by expert on the host
|
||||
pt.mark(kPfHostGroup, cs);
|
||||
// (the sync below also orders this layer's writes of slot/src/bounds after the previous
|
||||
// layer's kernels that read them)
|
||||
const bool grp_mapped = m.grp_host != nullptr;
|
||||
int32_t* ids_h = grp_mapped ? m.grp_host : m.ids_host.data();
|
||||
int32_t* slot_h = grp_mapped ? m.grp_host + m.grp_tk : m.slot_host.data();
|
||||
int32_t* src_h = grp_mapped ? m.grp_host + 2 * m.grp_tk : m.src_host.data();
|
||||
if (grp_mapped) copy_i32(m.grp_dev, m.ids, T * K, m.cs);
|
||||
else cudaMemcpyAsync(m.ids_host.data(), m.ids, (size_t) T * K * 4, cudaMemcpyDeviceToHost, m.cs);
|
||||
cudaStreamSynchronize(m.cs);
|
||||
pt.fold();
|
||||
std::fill(m.cnt.begin(), m.cnt.end(), 0);
|
||||
for (int64_t i = 0; i < T * K; ++i) {
|
||||
const int32_t e = ids_h[(size_t) i];
|
||||
if (e < 0 || e >= m.g->n_expert) { err = "prefill: routed id out of range"; return false; }
|
||||
++m.cnt[(size_t) e];
|
||||
}
|
||||
m.off[0] = 0;
|
||||
for (int64_t e = 0; e < m.g->n_expert; ++e) m.off[(size_t) e + 1] = m.off[(size_t) e] + m.cnt[(size_t) e];
|
||||
std::vector<int32_t> fill(m.off.begin(), m.off.end() - 1);
|
||||
for (int64_t i = 0; i < T * K; ++i) {
|
||||
const int32_t e = ids_h[(size_t) i];
|
||||
const int32_t p = fill[(size_t) e]++;
|
||||
slot_h[(size_t) i] = p;
|
||||
src_h[(size_t) p] = (int32_t) (i / K);
|
||||
}
|
||||
if (grp_mapped) {
|
||||
copy_i32(m.slot_dev, m.grp_dev + m.grp_tk, T * K, m.cs);
|
||||
copy_i32(m.src_dev, m.grp_dev + 2 * m.grp_tk, T * K, m.cs);
|
||||
} else {
|
||||
cudaMemcpyAsync(m.slot_dev, m.slot_host.data(), (size_t) T * K * 4, cudaMemcpyHostToDevice, m.cs);
|
||||
cudaMemcpyAsync(m.src_dev, m.src_host.data(), (size_t) T * K * 4, cudaMemcpyHostToDevice, m.cs);
|
||||
}
|
||||
// the experts, in id order: resident ones from VRAM, the others through the staging ring
|
||||
std::vector<int32_t> order;
|
||||
for (int32_t e = 0; e < m.g->n_expert; ++e) if (m.cnt[(size_t) e] > 0) order.push_back(e);
|
||||
// #136: STRATA_PF_FUSED=1 - the Q2_0 pack's experts on the fused int8 kernels (moe_fused.hpp),
|
||||
// grouped on the GPU: no host sync. Only where every expert's place is known before the routing -
|
||||
// the streamed walk, in which every non-resident expert of the layer comes through the ring in id
|
||||
// order - and where the MMQ buffers exist: they hold the fused path's own (the per-token int8
|
||||
// activations in Xq, the int8 H in H, the grouping tables in GU). Chunks below stream_all_min()
|
||||
// keep MMQ; without the variable nothing here runs. A native pack's layer takes the native kernels
|
||||
// (moe_fused_iq.hpp) where they cover its two formats, else MMQ (or the FP16 path: IQ1_M).
|
||||
const strata::kernels::cpu::ExpertLayout& lay = strata::kernels::cpu::expert_layout();
|
||||
const bool use_mmq = mmq_plan().any && mmq_plan().layer[(size_t) l];
|
||||
const int mmq_gt = lay.native ? lay.fmt[(size_t) l].gu_type : 42;
|
||||
const int mmq_dt = lay.native ? lay.fmt[(size_t) l].d_type : 42;
|
||||
const size_t mmq_gub = use_mmq ? mmq::matrix_bytes(mmq_gt, 1280, N) : 0;
|
||||
const size_t mmq_db = use_mmq ? mmq::matrix_bytes(mmq_dt, N, 640) : 0;
|
||||
pt.mark(kPfGather, cs);
|
||||
if (use_mmq) {
|
||||
// step 2b: the layer's activations as q8_1 rows in expert order, straight from `mixed`
|
||||
mmq::quantize(m.mixed, m.src_dev, m.Xq, mmq_gt, N, N, T * K, m.cs);
|
||||
// each group's rows: absolute bounds (gate/up reads the layer's rows), relative ones (down
|
||||
// reads the group's own quantized H)
|
||||
const size_t n = order.size(), ng = (n + MMQ_GROUP - 1) / MMQ_GROUP;
|
||||
m.bounds_host.resize(n + 1 + ng * (MMQ_GROUP + 1));
|
||||
for (size_t j = 0; j < n; ++j) m.bounds_host[j] = m.off[(size_t) order[j]];
|
||||
m.bounds_host[n] = (int32_t) (T * K);
|
||||
for (size_t g = 0; g < ng; ++g)
|
||||
for (size_t i = 0; i <= MMQ_GROUP; ++i)
|
||||
m.bounds_host[n + 1 + g * (MMQ_GROUP + 1) + i] =
|
||||
m.bounds_host[std::min(n, g * MMQ_GROUP + i)] - m.bounds_host[g * MMQ_GROUP];
|
||||
if (grp_mapped) {
|
||||
int32_t* bh = m.grp_host + 3 * m.grp_tk;
|
||||
std::memcpy(bh, m.bounds_host.data(), m.bounds_host.size() * 4);
|
||||
copy_i32(m.bounds_dev, m.grp_dev + 3 * m.grp_tk, (int64_t) m.bounds_host.size(), m.cs);
|
||||
} else {
|
||||
cudaMemcpyAsync(m.bounds_dev, m.bounds_host.data(), m.bounds_host.size() * 4,
|
||||
cudaMemcpyHostToDevice, m.cs);
|
||||
const bool fused_nat = use_mmq && stream_all && lay.native && fused::native_supported(mmq_gt, mmq_dt);
|
||||
const bool fused_l = (use_mmq && stream_all && !lay.native && fused::enabled()) || fused_nat;
|
||||
size_t n_order = 0; // the routed experts (the debug report; unknown when fused)
|
||||
if (fused_l) {
|
||||
if (static bool said = false; !said) {
|
||||
said = true;
|
||||
std::fprintf(stderr, "strata: prompt experts on the fused int8 kernels (STRATA_PF_FUSED=1, "
|
||||
"#136)\n");
|
||||
}
|
||||
} else {
|
||||
gather_rows16(m.mixed_h, m.src_dev, m.Xs, T * K, N, m.cs);
|
||||
}
|
||||
// Stage ahead: the copy stream moves blobs host -> device while the compute stream works.
|
||||
int stage_next = 0;
|
||||
std::vector<int> stage_of(order.size(), -1);
|
||||
// the unpinned ones are copied to pinned buffers by the stager's threads, in this order
|
||||
std::vector<int> job_of(order.size(), -1);
|
||||
if (!stream_all) {
|
||||
std::vector<Stager::Job> js;
|
||||
for (size_t j = 0; j < order.size(); ++j) {
|
||||
const int32_t e = order[j];
|
||||
if (m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0) continue;
|
||||
if (m.src->pinned(l, e)) continue;
|
||||
job_of[j] = (int) js.size();
|
||||
if (m.src->transient(l, e)) { // CS-T: copied by the source
|
||||
js.push_back({nullptr, (size_t) lay.blob_bytes(l), m.src, (int32_t) l, e});
|
||||
continue;
|
||||
}
|
||||
const uint8_t* b = m.src->blob(l, e);
|
||||
if (!b) { err = "prefill: expert source has no blob"; return false; }
|
||||
js.push_back({b, (size_t) lay.blob_bytes(l)});
|
||||
}
|
||||
m.stager->start(std::move(js));
|
||||
}
|
||||
StagerDone stager_done{stream_all ? nullptr : m.stager.get()};
|
||||
auto stage_one = [&](size_t j) -> bool {
|
||||
const int32_t e = order[j];
|
||||
const bool resident = m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0;
|
||||
if (resident) return true;
|
||||
const int sl = stage_next;
|
||||
stage_next = (stage_next + 1) % STAGE;
|
||||
const auto th = Clock::now();
|
||||
const bool pinned = m.src->pinned(l, e); // pinned: never transient
|
||||
const uint8_t* b = pinned ? m.src->blob(l, e) : nullptr;
|
||||
if (pinned && !b) { err = "prefill: expert source has no blob"; return false; }
|
||||
if (pinned) {
|
||||
// DMA straight from the page-locked arena: the copy stream only waits for the slot
|
||||
if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0);
|
||||
cudaMemcpyAsync(m.stage_dev[sl], b, (size_t) lay.blob_bytes(l), cudaMemcpyHostToDevice, m.copy);
|
||||
++stats_.experts_dma;
|
||||
} else {
|
||||
// copied to a pinned buffer by the stager (waits only if it is behind), then DMA
|
||||
const uint8_t* hb = m.stager->wait(job_of[j]);
|
||||
if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0);
|
||||
cudaMemcpyAsync(m.stage_dev[sl], hb, (size_t) lay.blob_bytes(l), cudaMemcpyHostToDevice, m.copy);
|
||||
m.stager->issued_one(job_of[j], m.copy);
|
||||
}
|
||||
cudaEventRecord(m.copied[sl], m.copy);
|
||||
m.stage_live[sl] = true;
|
||||
stage_of[j] = sl;
|
||||
stats_.ms_experts_host += ms_since(th);
|
||||
++stats_.experts_streamed;
|
||||
return true;
|
||||
};
|
||||
// one expert's products from its blob on the device; `slot` (a ring slot, or -1 for a resident
|
||||
// expert) is released once the blob is read
|
||||
auto compute = [&](size_t j, const uint8_t* blob_dev, int slot) -> bool {
|
||||
const int32_t e = order[j];
|
||||
pt.mark(kPfDequant, cs);
|
||||
if (use_mmq) {
|
||||
// gather the expert into its group slot (GGUF blocks, unchanged or converted)
|
||||
const size_t q = j % MMQ_GROUP;
|
||||
if (lay.native) {
|
||||
const auto& f = lay.fmt[(size_t) l];
|
||||
mmq::gather_native(blob_dev, blob_dev + f.up_off, mmq_gub / 2, blob_dev + f.down_off,
|
||||
mmq_db, m.grp_gu + q * mmq_gub, m.grp_d + q * mmq_db, m.cs);
|
||||
} else {
|
||||
mmq::gather_strata_q2(blob_dev, m.grp_gu + q * mmq_gub, m.grp_d + q * mmq_db, m.cs);
|
||||
}
|
||||
if (slot >= 0) cudaEventRecord(m.used[slot], m.cs);
|
||||
if (q + 1 < MMQ_GROUP && j + 1 < order.size()) return true;
|
||||
// the group's products: gate/up, swiglu, the group's H to q8_1, down
|
||||
const size_t j0 = j - q, g = j0 / MMQ_GROUP, n = order.size();
|
||||
const int ngx = (int) (q + 1);
|
||||
const int64_t r0 = m.bounds_host[j0], nr = m.bounds_host[j + 1] - r0;
|
||||
int64_t maxr = 0;
|
||||
for (size_t i = j0; i <= j; ++i) maxr = std::max<int64_t>(maxr, m.cnt[(size_t) order[i]]);
|
||||
pt.mark(kPfGemmGU, cs);
|
||||
// the zeroed tail after the group's last expert (see MMQ_TAIL)
|
||||
cudaMemsetAsync(m.grp_gu + (size_t) ngx * mmq_gub, 0, MMQ_TAIL, m.cs);
|
||||
cudaMemsetAsync(m.grp_d + (size_t) ngx * mmq_db, 0, MMQ_TAIL, m.cs);
|
||||
mmq::Product gu;
|
||||
gu.w = m.grp_gu; gu.type = mmq_gt; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = mmq_gub;
|
||||
gu.n = ngx; gu.xq = m.Xq; gu.bounds = m.bounds_dev + j0; gu.ids = m.ids_identity;
|
||||
gu.total_rows = T * K; gu.max_rows = maxr; gu.dst = m.GU; gu.ld_dst = 1280;
|
||||
m.mmq_ctx->run(gu, m.cs);
|
||||
mmq::swiglu(m.GU + r0 * 1280, m.H + r0 * 640, nr, 640, !lay.native, m.cs);
|
||||
pt.mark(kPfGemmD, cs);
|
||||
mmq::quantize(m.H + r0 * 640, nullptr, m.Hq, mmq_dt, 640, 640, nr, m.cs);
|
||||
mmq::Product dn;
|
||||
dn.w = m.grp_d; dn.type = mmq_dt; dn.w_rows = N; dn.w_cols = 640; dn.expert_bytes = mmq_db;
|
||||
dn.n = ngx; dn.xq = m.Hq; dn.bounds = m.bounds_dev + n + 1 + g * (MMQ_GROUP + 1);
|
||||
dn.ids = m.ids_identity; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = m.Dm + r0 * N;
|
||||
dn.ld_dst = N;
|
||||
m.mmq_ctx->run(dn, m.cs);
|
||||
return true;
|
||||
}
|
||||
const int q = (int) (j % DQ);
|
||||
if (lay.native) {
|
||||
// plan v0.3 P6: a native pack's layer, dequantized by llama.cpp's own formulas
|
||||
const auto& f = lay.fmt[(size_t) l];
|
||||
strata::kernels::iq_dequant_gu_f16(f.gu_type, blob_dev, blob_dev + f.up_off, f.n_ff, f.n_embd,
|
||||
m.dq_gu[q], m.cs);
|
||||
strata::kernels::iq_dequant_f16(f.d_type, blob_dev + f.down_off, f.n_embd * f.n_ff, m.dq_d[q], m.cs);
|
||||
} else {
|
||||
blob_dequant_f16(blob_dev, m.dq_gu[q], m.dq_d[q], m.cs);
|
||||
}
|
||||
if (slot >= 0) cudaEventRecord(m.used[slot], m.cs);
|
||||
const int64_t o0 = m.off[(size_t) e], ne = m.cnt[(size_t) e];
|
||||
pt.mark(kPfGemmGU, cs);
|
||||
m.gemm.f16(m.Xs + o0 * N, m.dq_gu[q], m.GU + o0 * 1280, ne, 1280, N);
|
||||
swiglu_interleaved(m.GU + o0 * 1280, m.Hh + o0 * 640, ne, m.cs);
|
||||
pt.mark(kPfGemmD, cs);
|
||||
m.gemm.f16(m.Hh + o0 * 640, m.dq_d[q], m.Dm + o0 * N, ne, N, 640);
|
||||
return true;
|
||||
};
|
||||
if (!stream_all) {
|
||||
size_t staged = 0;
|
||||
const size_t lookahead = STAGE - 1;
|
||||
for (size_t j = 0; j < order.size(); ++j) {
|
||||
while (staged < order.size() && staged <= j + lookahead) {
|
||||
if (!stage_one(staged)) return false;
|
||||
++staged;
|
||||
}
|
||||
const int32_t e = order[j];
|
||||
if (stage_of[j] < 0) {
|
||||
++stats_.experts_resident;
|
||||
if (!compute(j, m.cache->device_slot(m.host_res[(size_t) l * m.g->n_expert + e]), -1)) return false;
|
||||
} else {
|
||||
pt.mark(kPfWaitCopy, cs);
|
||||
cudaStreamWaitEvent(m.cs, m.copied[stage_of[j]], 0);
|
||||
if (!compute(j, m.stage_dev[stage_of[j]], stage_of[j])) return false;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// the streamed walk: this layer's entries [k, kend) in id order; an entry the routing did not
|
||||
// pick only gives its slot back
|
||||
pt.mark(kPfGather, cs);
|
||||
// the layer's input to int8 once per token; the rows of each expert from the router's ids
|
||||
if (fused_nat) fused::quantize_act_native(m.mixed, T, N, m.Xq, m.cs);
|
||||
else fused::quantize_act(m.mixed, T, N, m.Xq, m.cs);
|
||||
fused::group(m.ids, T * K, (int) K, (int) m.g->n_expert, m.GU, m.slot_dev, m.src_dev, m.cs);
|
||||
// launches over the experts in id order, each at most kMaxBatch experts of which at most a
|
||||
// third of the ring streamed: the next batch's blobs arrive while one computes. An expert
|
||||
// the routing did not pick has no tiles; its ring slot is given back with its batch.
|
||||
const size_t per = (size_t) std::max(1, m.ring / 3);
|
||||
size_t k = seq_start[(size_t) l];
|
||||
const size_t kend = seq_start[(size_t) l + 1];
|
||||
auto release_to = [&](int32_t e_stop) {
|
||||
while (k < kend && seq[k].e < e_stop) {
|
||||
cudaEventRecord(m.used[k % (size_t) m.ring], m.cs);
|
||||
consumed = ++k;
|
||||
give_back(consumed);
|
||||
for (int32_t e = 0; e < m.g->n_expert;) {
|
||||
fused::Batch b;
|
||||
b.e0 = e;
|
||||
const size_t k0 = k;
|
||||
for (; e < m.g->n_expert && e - b.e0 < fused::kMaxBatch; ++e) {
|
||||
if (k < kend && seq[k].e == e) {
|
||||
if (k - k0 == per) break;
|
||||
b.blob[e - b.e0] = m.stage_dev[k % (size_t) m.ring];
|
||||
++k;
|
||||
} else { // not streamed: resident (the walk streams all others)
|
||||
b.blob[e - b.e0] = m.cache->device_slot(m.host_res[(size_t) l * m.g->n_expert + e]);
|
||||
// counted whether routed or not (the routing stays on the GPU): at a streamed
|
||||
// chunk's size (>= 1024 tokens x 10 of 512) nearly every expert is routed
|
||||
++stats_.experts_resident;
|
||||
}
|
||||
}
|
||||
};
|
||||
for (size_t j = 0; j < order.size(); ++j) {
|
||||
const int32_t e = order[j];
|
||||
release_to(e);
|
||||
if (k < kend && seq[k].e == e) {
|
||||
const int sl = (int) (k % (size_t) m.ring);
|
||||
b.e1 = e;
|
||||
if (k > k0) {
|
||||
// one copy stream, in order: the batch's last blob covers the others
|
||||
pt.mark(kPfWaitCopy, cs);
|
||||
wait_issued(k);
|
||||
cudaStreamWaitEvent(m.cs, m.copied[sl], 0);
|
||||
if (!compute(j, m.stage_dev[sl], sl)) return false;
|
||||
consumed = ++k;
|
||||
give_back(consumed);
|
||||
wait_issued(k - 1);
|
||||
cudaStreamWaitEvent(m.cs, m.copied[(k - 1) % (size_t) m.ring], 0);
|
||||
}
|
||||
pt.mark(kPfGemmGU, cs);
|
||||
if (fused_nat) {
|
||||
const auto& f = lay.fmt[(size_t) l];
|
||||
const fused::NativeGeom ng{f.gu_type, f.d_type, f.gu_row, f.d_row, f.up_off, f.down_off};
|
||||
fused::experts_native(b, ng, (int) m.g->n_expert, T * K, m.GU, m.Xq, m.src_dev, m.H,
|
||||
m.Dm, m.cs);
|
||||
} else {
|
||||
++stats_.experts_resident;
|
||||
if (!compute(j, m.cache->device_slot(m.host_res[(size_t) l * m.g->n_expert + e]), -1)) return false;
|
||||
fused::experts(b, (int) m.g->n_expert, T * K, m.GU, m.Xq, m.src_dev, m.H, m.Dm, m.cs);
|
||||
}
|
||||
for (size_t kk = k0; kk < k; ++kk) cudaEventRecord(m.used[kk % (size_t) m.ring], m.cs);
|
||||
if (k > k0) {
|
||||
consumed = k;
|
||||
give_back(consumed);
|
||||
}
|
||||
}
|
||||
release_to(m.g->n_expert);
|
||||
} else {
|
||||
// group the (token, k) pairs by expert on the host
|
||||
pt.mark(kPfHostGroup, cs);
|
||||
// (the sync below also orders this layer's writes of slot/src/bounds after the previous
|
||||
// layer's kernels that read them)
|
||||
const bool grp_mapped = m.grp_host != nullptr;
|
||||
int32_t* ids_h = grp_mapped ? m.grp_host : m.ids_host.data();
|
||||
int32_t* slot_h = grp_mapped ? m.grp_host + m.grp_tk : m.slot_host.data();
|
||||
int32_t* src_h = grp_mapped ? m.grp_host + 2 * m.grp_tk : m.src_host.data();
|
||||
if (grp_mapped) copy_i32(m.grp_dev, m.ids, T * K, m.cs);
|
||||
else cudaMemcpyAsync(m.ids_host.data(), m.ids, (size_t) T * K * 4, cudaMemcpyDeviceToHost, m.cs);
|
||||
cudaStreamSynchronize(m.cs);
|
||||
pt.fold();
|
||||
std::fill(m.cnt.begin(), m.cnt.end(), 0);
|
||||
for (int64_t i = 0; i < T * K; ++i) {
|
||||
const int32_t e = ids_h[(size_t) i];
|
||||
if (e < 0 || e >= m.g->n_expert) { err = "prefill: routed id out of range"; return false; }
|
||||
++m.cnt[(size_t) e];
|
||||
}
|
||||
m.off[0] = 0;
|
||||
for (int64_t e = 0; e < m.g->n_expert; ++e) m.off[(size_t) e + 1] = m.off[(size_t) e] + m.cnt[(size_t) e];
|
||||
std::vector<int32_t> fill(m.off.begin(), m.off.end() - 1);
|
||||
for (int64_t i = 0; i < T * K; ++i) {
|
||||
const int32_t e = ids_h[(size_t) i];
|
||||
const int32_t p = fill[(size_t) e]++;
|
||||
slot_h[(size_t) i] = p;
|
||||
src_h[(size_t) p] = (int32_t) (i / K);
|
||||
}
|
||||
if (grp_mapped) {
|
||||
copy_i32(m.slot_dev, m.grp_dev + m.grp_tk, T * K, m.cs);
|
||||
copy_i32(m.src_dev, m.grp_dev + 2 * m.grp_tk, T * K, m.cs);
|
||||
} else {
|
||||
cudaMemcpyAsync(m.slot_dev, m.slot_host.data(), (size_t) T * K * 4, cudaMemcpyHostToDevice, m.cs);
|
||||
cudaMemcpyAsync(m.src_dev, m.src_host.data(), (size_t) T * K * 4, cudaMemcpyHostToDevice, m.cs);
|
||||
}
|
||||
// the experts, in id order: resident ones from VRAM, the others through the staging ring
|
||||
std::vector<int32_t> order;
|
||||
for (int32_t e = 0; e < m.g->n_expert; ++e) if (m.cnt[(size_t) e] > 0) order.push_back(e);
|
||||
n_order = order.size();
|
||||
const size_t mmq_gub = use_mmq ? mmq::matrix_bytes(mmq_gt, 1280, N) : 0;
|
||||
const size_t mmq_db = use_mmq ? mmq::matrix_bytes(mmq_dt, N, 640) : 0;
|
||||
pt.mark(kPfGather, cs);
|
||||
if (use_mmq) {
|
||||
// step 2b: the layer's activations as q8_1 rows in expert order, straight from `mixed`
|
||||
mmq::quantize(m.mixed, m.src_dev, m.Xq, mmq_gt, N, N, T * K, m.cs);
|
||||
// each group's rows: absolute bounds (gate/up reads the layer's rows), relative ones (down
|
||||
// reads the group's own quantized H)
|
||||
const size_t n = order.size(), ng = (n + MMQ_GROUP - 1) / MMQ_GROUP;
|
||||
m.bounds_host.resize(n + 1 + ng * (MMQ_GROUP + 1));
|
||||
for (size_t j = 0; j < n; ++j) m.bounds_host[j] = m.off[(size_t) order[j]];
|
||||
m.bounds_host[n] = (int32_t) (T * K);
|
||||
for (size_t g = 0; g < ng; ++g)
|
||||
for (size_t i = 0; i <= MMQ_GROUP; ++i)
|
||||
m.bounds_host[n + 1 + g * (MMQ_GROUP + 1) + i] =
|
||||
m.bounds_host[std::min(n, g * MMQ_GROUP + i)] - m.bounds_host[g * MMQ_GROUP];
|
||||
if (grp_mapped) {
|
||||
int32_t* bh = m.grp_host + 3 * m.grp_tk;
|
||||
std::memcpy(bh, m.bounds_host.data(), m.bounds_host.size() * 4);
|
||||
copy_i32(m.bounds_dev, m.grp_dev + 3 * m.grp_tk, (int64_t) m.bounds_host.size(), m.cs);
|
||||
} else {
|
||||
cudaMemcpyAsync(m.bounds_dev, m.bounds_host.data(), m.bounds_host.size() * 4,
|
||||
cudaMemcpyHostToDevice, m.cs);
|
||||
}
|
||||
} else {
|
||||
gather_rows16(m.mixed_h, m.src_dev, m.Xs, T * K, N, m.cs);
|
||||
}
|
||||
// Stage ahead: the copy stream moves blobs host -> device while the compute stream works.
|
||||
int stage_next = 0;
|
||||
std::vector<int> stage_of(order.size(), -1);
|
||||
// the unpinned ones are copied to pinned buffers by the stager's threads, in this order
|
||||
std::vector<int> job_of(order.size(), -1);
|
||||
if (!stream_all) {
|
||||
std::vector<Stager::Job> js;
|
||||
for (size_t j = 0; j < order.size(); ++j) {
|
||||
const int32_t e = order[j];
|
||||
if (m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0) continue;
|
||||
if (m.src->pinned(l, e)) continue;
|
||||
job_of[j] = (int) js.size();
|
||||
if (m.src->transient(l, e)) { // CS-T: copied by the source
|
||||
js.push_back({nullptr, (size_t) lay.blob_bytes(l), m.src, (int32_t) l, e});
|
||||
continue;
|
||||
}
|
||||
const uint8_t* b = m.src->blob(l, e);
|
||||
if (!b) { err = "prefill: expert source has no blob"; return false; }
|
||||
js.push_back({b, (size_t) lay.blob_bytes(l)});
|
||||
}
|
||||
m.stager->start(std::move(js));
|
||||
}
|
||||
StagerDone stager_done{stream_all ? nullptr : m.stager.get()};
|
||||
auto stage_one = [&](size_t j) -> bool {
|
||||
const int32_t e = order[j];
|
||||
const bool resident = m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0;
|
||||
if (resident) return true;
|
||||
const int sl = stage_next;
|
||||
stage_next = (stage_next + 1) % STAGE;
|
||||
const auto th = Clock::now();
|
||||
const bool pinned = m.src->pinned(l, e); // pinned: never transient
|
||||
const uint8_t* b = pinned ? m.src->blob(l, e) : nullptr;
|
||||
if (pinned && !b) { err = "prefill: expert source has no blob"; return false; }
|
||||
if (pinned) {
|
||||
// DMA straight from the page-locked arena: the copy stream only waits for the slot
|
||||
if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0);
|
||||
cudaMemcpyAsync(m.stage_dev[sl], b, (size_t) lay.blob_bytes(l), cudaMemcpyHostToDevice, m.copy);
|
||||
++stats_.experts_dma;
|
||||
} else {
|
||||
// copied to a pinned buffer by the stager (waits only if it is behind), then DMA
|
||||
const uint8_t* hb = m.stager->wait(job_of[j]);
|
||||
if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0);
|
||||
cudaMemcpyAsync(m.stage_dev[sl], hb, (size_t) lay.blob_bytes(l), cudaMemcpyHostToDevice, m.copy);
|
||||
m.stager->issued_one(job_of[j], m.copy);
|
||||
}
|
||||
cudaEventRecord(m.copied[sl], m.copy);
|
||||
m.stage_live[sl] = true;
|
||||
stage_of[j] = sl;
|
||||
stats_.ms_experts_host += ms_since(th);
|
||||
++stats_.experts_streamed;
|
||||
return true;
|
||||
};
|
||||
// one expert's products from its blob on the device; `slot` (a ring slot, or -1 for a resident
|
||||
// expert) is released once the blob is read
|
||||
auto compute = [&](size_t j, const uint8_t* blob_dev, int slot) -> bool {
|
||||
const int32_t e = order[j];
|
||||
pt.mark(kPfDequant, cs);
|
||||
if (use_mmq) {
|
||||
// gather the expert into its group slot (GGUF blocks, unchanged or converted)
|
||||
const size_t q = j % MMQ_GROUP;
|
||||
if (lay.native) {
|
||||
const auto& f = lay.fmt[(size_t) l];
|
||||
mmq::gather_native(blob_dev, blob_dev + f.up_off, mmq_gub / 2, blob_dev + f.down_off,
|
||||
mmq_db, m.grp_gu + q * mmq_gub, m.grp_d + q * mmq_db, m.cs);
|
||||
} else {
|
||||
mmq::gather_strata_q2(blob_dev, m.grp_gu + q * mmq_gub, m.grp_d + q * mmq_db, m.cs);
|
||||
}
|
||||
if (slot >= 0) cudaEventRecord(m.used[slot], m.cs);
|
||||
if (q + 1 < MMQ_GROUP && j + 1 < order.size()) return true;
|
||||
// the group's products: gate/up, swiglu, the group's H to q8_1, down
|
||||
const size_t j0 = j - q, g = j0 / MMQ_GROUP, n = order.size();
|
||||
const int ngx = (int) (q + 1);
|
||||
const int64_t r0 = m.bounds_host[j0], nr = m.bounds_host[j + 1] - r0;
|
||||
int64_t maxr = 0;
|
||||
for (size_t i = j0; i <= j; ++i) maxr = std::max<int64_t>(maxr, m.cnt[(size_t) order[i]]);
|
||||
pt.mark(kPfGemmGU, cs);
|
||||
// the zeroed tail after the group's last expert (see MMQ_TAIL)
|
||||
cudaMemsetAsync(m.grp_gu + (size_t) ngx * mmq_gub, 0, MMQ_TAIL, m.cs);
|
||||
cudaMemsetAsync(m.grp_d + (size_t) ngx * mmq_db, 0, MMQ_TAIL, m.cs);
|
||||
mmq::Product gu;
|
||||
gu.w = m.grp_gu; gu.type = mmq_gt; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = mmq_gub;
|
||||
gu.n = ngx; gu.xq = m.Xq; gu.bounds = m.bounds_dev + j0; gu.ids = m.ids_identity;
|
||||
gu.total_rows = T * K; gu.max_rows = maxr; gu.dst = m.GU; gu.ld_dst = 1280;
|
||||
m.mmq_ctx->run(gu, m.cs);
|
||||
mmq::swiglu(m.GU + r0 * 1280, m.H + r0 * 640, nr, 640, !lay.native, m.cs);
|
||||
pt.mark(kPfGemmD, cs);
|
||||
mmq::quantize(m.H + r0 * 640, nullptr, m.Hq, mmq_dt, 640, 640, nr, m.cs);
|
||||
mmq::Product dn;
|
||||
dn.w = m.grp_d; dn.type = mmq_dt; dn.w_rows = N; dn.w_cols = 640; dn.expert_bytes = mmq_db;
|
||||
dn.n = ngx; dn.xq = m.Hq; dn.bounds = m.bounds_dev + n + 1 + g * (MMQ_GROUP + 1);
|
||||
dn.ids = m.ids_identity; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = m.Dm + r0 * N;
|
||||
dn.ld_dst = N;
|
||||
m.mmq_ctx->run(dn, m.cs);
|
||||
return true;
|
||||
}
|
||||
const int q = (int) (j % DQ);
|
||||
if (lay.native) {
|
||||
// plan v0.3 P6: a native pack's layer, dequantized by llama.cpp's own formulas
|
||||
const auto& f = lay.fmt[(size_t) l];
|
||||
strata::kernels::iq_dequant_gu_f16(f.gu_type, blob_dev, blob_dev + f.up_off, f.n_ff, f.n_embd,
|
||||
m.dq_gu[q], m.cs);
|
||||
strata::kernels::iq_dequant_f16(f.d_type, blob_dev + f.down_off, f.n_embd * f.n_ff, m.dq_d[q], m.cs);
|
||||
} else {
|
||||
blob_dequant_f16(blob_dev, m.dq_gu[q], m.dq_d[q], m.cs);
|
||||
}
|
||||
if (slot >= 0) cudaEventRecord(m.used[slot], m.cs);
|
||||
const int64_t o0 = m.off[(size_t) e], ne = m.cnt[(size_t) e];
|
||||
pt.mark(kPfGemmGU, cs);
|
||||
m.gemm.f16(m.Xs + o0 * N, m.dq_gu[q], m.GU + o0 * 1280, ne, 1280, N);
|
||||
swiglu_interleaved(m.GU + o0 * 1280, m.Hh + o0 * 640, ne, m.cs);
|
||||
pt.mark(kPfGemmD, cs);
|
||||
m.gemm.f16(m.Hh + o0 * 640, m.dq_d[q], m.Dm + o0 * N, ne, N, 640);
|
||||
return true;
|
||||
};
|
||||
if (!stream_all) {
|
||||
size_t staged = 0;
|
||||
const size_t lookahead = STAGE - 1;
|
||||
for (size_t j = 0; j < order.size(); ++j) {
|
||||
while (staged < order.size() && staged <= j + lookahead) {
|
||||
if (!stage_one(staged)) return false;
|
||||
++staged;
|
||||
}
|
||||
const int32_t e = order[j];
|
||||
if (stage_of[j] < 0) {
|
||||
++stats_.experts_resident;
|
||||
if (!compute(j, m.cache->device_slot(m.host_res[(size_t) l * m.g->n_expert + e]), -1)) return false;
|
||||
} else {
|
||||
pt.mark(kPfWaitCopy, cs);
|
||||
cudaStreamWaitEvent(m.cs, m.copied[stage_of[j]], 0);
|
||||
if (!compute(j, m.stage_dev[stage_of[j]], stage_of[j])) return false;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// the streamed walk: this layer's entries [k, kend) in id order; an entry the routing did not
|
||||
// pick only gives its slot back
|
||||
size_t k = seq_start[(size_t) l];
|
||||
const size_t kend = seq_start[(size_t) l + 1];
|
||||
auto release_to = [&](int32_t e_stop) {
|
||||
while (k < kend && seq[k].e < e_stop) {
|
||||
cudaEventRecord(m.used[k % (size_t) m.ring], m.cs);
|
||||
consumed = ++k;
|
||||
give_back(consumed);
|
||||
}
|
||||
};
|
||||
for (size_t j = 0; j < order.size(); ++j) {
|
||||
const int32_t e = order[j];
|
||||
release_to(e);
|
||||
if (k < kend && seq[k].e == e) {
|
||||
const int sl = (int) (k % (size_t) m.ring);
|
||||
pt.mark(kPfWaitCopy, cs);
|
||||
wait_issued(k);
|
||||
cudaStreamWaitEvent(m.cs, m.copied[sl], 0);
|
||||
if (!compute(j, m.stage_dev[sl], sl)) return false;
|
||||
consumed = ++k;
|
||||
give_back(consumed);
|
||||
} else {
|
||||
++stats_.experts_resident;
|
||||
if (!compute(j, m.cache->device_slot(m.host_res[(size_t) l * m.g->n_expert + e]), -1)) return false;
|
||||
}
|
||||
}
|
||||
release_to(m.g->n_expert);
|
||||
}
|
||||
}
|
||||
pt.mark(kPfCombine, cs);
|
||||
moe_combine(m.Dm, m.slot_dev, m.w, m.shared, m.sg, m.bo, T, m.cs);
|
||||
@@ -1821,14 +1955,16 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
|
||||
for (float v : h) c += !std::isfinite(v);
|
||||
return c;
|
||||
};
|
||||
const int64_t bgu = bad(m.GU, T * K * 1280), bdm = bad(m.Dm, T * K * N), bbo = bad(m.bo, T * N);
|
||||
const int64_t bh = m.H ? bad(m.H, T * K * 640) : -1;
|
||||
// (fused: GU and H hold the grouping tables and int8 rows, not floats)
|
||||
const int64_t bgu = fused_l ? 0 : bad(m.GU, T * K * 1280), bdm = bad(m.Dm, T * K * N),
|
||||
bbo = bad(m.bo, T * N);
|
||||
const int64_t bh = m.H && !fused_l ? bad(m.H, T * K * 640) : -1;
|
||||
static int64_t reported = -1;
|
||||
if ((bgu || bdm || bbo || bh > 0) && reported != stats_.chunks) {
|
||||
reported = stats_.chunks;
|
||||
std::fprintf(stderr, "strata dbg: layer %lld (mmq %d, types %d/%d, %zu experts): non-finite GU %lld "
|
||||
"H %lld Dm %lld bo %lld of T %lld\n", (long long) l, (int) use_mmq, mmq_gt, mmq_dt,
|
||||
order.size(), (long long) bgu, (long long) bh, (long long) bdm, (long long) bbo,
|
||||
n_order, (long long) bgu, (long long) bh, (long long) bdm, (long long) bbo,
|
||||
(long long) T);
|
||||
}
|
||||
}
|
||||
|
||||
+230
-49
@@ -303,6 +303,11 @@ struct Options {
|
||||
/// trace is what the plan's `h = 0.6447` refers to, and compulsory-miss measured 0.4864 because it fills
|
||||
/// with whatever the prompt touched FIRST. Empty means no profile.
|
||||
std::string expert_profile;
|
||||
/// #477 (--serve, opt-in): where to save what the adaptive tier learned, as a profile `--expert-profile` reads
|
||||
/// (the resident experts first, then the routing counted since the start); on QUIT and every
|
||||
/// `expert_profile_save_min` minutes between requests. Empty (the default): nothing is counted or written.
|
||||
std::string expert_profile_save;
|
||||
double expert_profile_save_min = 10.0;
|
||||
/// R4.2d: **ON by default**, because the measurement is unambiguous and the alternative is known-broken.
|
||||
/// Without it, 17 of 10,562 layers had the hit work done when the pool returned; with it, 9,190. The
|
||||
/// A/B arm is `--no-hit-poke`.
|
||||
@@ -530,6 +535,10 @@ void usage() {
|
||||
" real graph that --stage-timing cannot give. Prints and exits.\n"
|
||||
" --expert-profile P R4.2e: pre-load the VRAM tier from a `profile.bin` (see\n"
|
||||
" tools/make_profile.py) instead of admitting on first use.\n"
|
||||
" --expert-profile-save P --serve, #477: save what the adaptive tier learned (the experts in\n"
|
||||
" VRAM, then the routing counted since the start) as a profile at P, on\n"
|
||||
" QUIT and every --expert-profile-save-every MIN minutes (default 10;\n"
|
||||
" 0 = on QUIT only) between requests; start from it with --expert-profile P\n"
|
||||
" --no-hit-poke R4.2d's A/B arm. The hit path pokes the driver once right after its\n"
|
||||
" launch so the GPU starts while the CPU pool runs; without it the work\n"
|
||||
" waits for the next driver entry and does not overlap at all.\n"
|
||||
@@ -947,35 +956,81 @@ bool load_control_vectors(const Options& o, const strata::core::ModelGeometry& g
|
||||
// The effective host->device bandwidth of the PCIe link: copies from pinned host memory, as the expert arena's
|
||||
// reads are. The native default share (0.55) was measured on x16 links (~26-28 GB/s); a x8 card in a x8 slot
|
||||
// carries about half of that. Returns < 0 when the probe cannot run (then the caller keeps the default).
|
||||
double probe_pcie_h2d_gbps() {
|
||||
//
|
||||
// #485: a single timed burst read an x16 PCIe 4 link (RTX A3000 laptop) at 18.5, 6.9 and 5.8 GB/s in three starts -
|
||||
// a link still in a low-power state, other DMA, a context not yet up to speed - and the low reading set the share.
|
||||
// Such a disturbance only ever slows a copy: nothing makes a correctly timed copy faster than the link carries. So
|
||||
// the same 1 GiB is now copied as four bursts of 256 MiB, timed one by one, and the fastest is the link's figure (the
|
||||
// median would still follow a disturbance that lasts through half the bursts). Each burst takes ~10 ms on an x16
|
||||
// PCIe 4 link, so the probe takes no longer than the one 1 GiB burst did. `samples`, when given, gets every
|
||||
// burst's reading for the log.
|
||||
double probe_pcie_h2d_gbps(std::string* samples = nullptr) {
|
||||
constexpr size_t kBytes = 256ull << 20;
|
||||
constexpr int kIters = 4;
|
||||
constexpr int kBursts = 4;
|
||||
void* h = nullptr;
|
||||
void* d = nullptr;
|
||||
cudaEvent_t ev0, ev1;
|
||||
if (cudaMallocHost(&h, kBytes) != cudaSuccess) return -1.0;
|
||||
if (cudaMalloc(&d, kBytes) != cudaSuccess || cudaEventCreate(&ev0) != cudaSuccess ||
|
||||
cudaEventCreate(&ev1) != cudaSuccess) {
|
||||
if (d != nullptr) cudaFree(d);
|
||||
if (cudaMalloc(&d, kBytes) != cudaSuccess) {
|
||||
cudaFreeHost(h);
|
||||
return -1.0;
|
||||
}
|
||||
std::memset(h, 0, kBytes); // fault the pages in before timing
|
||||
cudaMemcpyAsync(d, h, kBytes, cudaMemcpyHostToDevice); // warmup: context up, copy engine primed
|
||||
cudaEventRecord(ev0);
|
||||
for (int i = 0; i < kIters; ++i) cudaMemcpyAsync(d, h, kBytes, cudaMemcpyHostToDevice);
|
||||
cudaEventRecord(ev1);
|
||||
const bool ok = cudaEventSynchronize(ev1) == cudaSuccess;
|
||||
float ms = 0.f;
|
||||
const bool timed = ok && cudaEventElapsedTime(&ms, ev0, ev1) == cudaSuccess && ms > 0.01f;
|
||||
const double bw = timed ? ((double) kIters * (double) kBytes / (ms * 1e-3)) / 1e9 : -1.0;
|
||||
cudaEventDestroy(ev0);
|
||||
cudaEventDestroy(ev1);
|
||||
float ms[kBursts] = {};
|
||||
#if defined(STRATA_USE_HIP) && defined(_WIN32)
|
||||
// Windows HIP: the events do not bracket the copies there (an RX 6800 read 3,300-26,000 GB/s, so every link kept
|
||||
// the x16 share), so each burst is timed on the host between two synchronizes: 256 MiB takes ~10 ms, so the
|
||||
// synchronize around it hardly matters
|
||||
bool ok = cudaDeviceSynchronize() == cudaSuccess;
|
||||
for (int b = 0; b < kBursts && ok; ++b) {
|
||||
const Clock::time_point t0 = Clock::now();
|
||||
cudaMemcpyAsync(d, h, kBytes, cudaMemcpyHostToDevice);
|
||||
ok = cudaDeviceSynchronize() == cudaSuccess;
|
||||
ms[b] = (float) std::chrono::duration<double, std::milli>(Clock::now() - t0).count();
|
||||
}
|
||||
#else
|
||||
// the bursts run back to back on the stream, an event between each two
|
||||
cudaEvent_t ev[kBursts + 1];
|
||||
int n_ev = 0;
|
||||
while (n_ev <= kBursts && cudaEventCreate(&ev[n_ev]) == cudaSuccess) ++n_ev;
|
||||
bool ok = n_ev == kBursts + 1;
|
||||
if (ok) {
|
||||
cudaEventRecord(ev[0]);
|
||||
for (int b = 0; b < kBursts; ++b) {
|
||||
cudaMemcpyAsync(d, h, kBytes, cudaMemcpyHostToDevice);
|
||||
cudaEventRecord(ev[b + 1]);
|
||||
}
|
||||
ok = cudaEventSynchronize(ev[kBursts]) == cudaSuccess;
|
||||
for (int b = 0; b < kBursts && ok; ++b) ok = cudaEventElapsedTime(&ms[b], ev[b], ev[b + 1]) == cudaSuccess;
|
||||
}
|
||||
for (int i = 0; i < n_ev; ++i) cudaEventDestroy(ev[i]);
|
||||
#endif
|
||||
double bw = -1.0;
|
||||
if (samples != nullptr) samples->clear();
|
||||
for (int b = 0; b < kBursts && ok; ++b) {
|
||||
const double s = ms[b] > 0.01f ? (double) kBytes / (ms[b] * 1e-3) / 1e9 : -1.0;
|
||||
bw = std::max(bw, s);
|
||||
if (samples != nullptr) {
|
||||
char buf[32];
|
||||
std::snprintf(buf, sizeof(buf), "%s%.1f", b == 0 ? "" : " ", s);
|
||||
*samples += buf;
|
||||
}
|
||||
}
|
||||
if (!ok) (void) cudaGetLastError();
|
||||
cudaFree(d);
|
||||
cudaFreeHost(h);
|
||||
return bw;
|
||||
}
|
||||
|
||||
// The PCIe share of the missed experts for a link measured at `gbps`: `base` (the share measured on x16 links) from
|
||||
// 20 GB/s up, and below that in proportion to the bandwidth, so the time the link spends on its share stays about
|
||||
// what the x16 share costs. Continuous (#485): before, 19.9 GB/s gave 0.42 and 20.0 the full 0.55 (and 4.0 GB/s
|
||||
// gave 0.08, 3.9 none), so a reading near either edge moved the share by a quarter of its range. From 20 GB/s up
|
||||
// (an x16 PCIe 4/5 link: ~26-28 GB/s) the share is unchanged.
|
||||
double pcie_frac_for_gbps(double gbps, double base) {
|
||||
return gbps <= 0.0 ? base : base * std::min(1.0, gbps / 20.0);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
@@ -1210,6 +1265,9 @@ int main(int argc, char** argv) {
|
||||
else if (a == "--expert-cache-per-layer") o.expert_cache_per_layer = true;
|
||||
else if (a == "--no-hit-poke") o.no_hit_poke = true;
|
||||
else if (a == "--expert-profile") o.expert_profile = next("--expert-profile");
|
||||
else if (a == "--expert-profile-save") o.expert_profile_save = next("--expert-profile-save");
|
||||
else if (a == "--expert-profile-save-every")
|
||||
o.expert_profile_save_min = std::atof(next("--expert-profile-save-every"));
|
||||
else if (a == "--gpu-stages") o.gpu_stages = true;
|
||||
else if (a == "--mmap-experts") o.mmap_experts = true;
|
||||
else if (a == "--shared-expert-arena") o.shared_expert_arena = next("--shared-expert-arena");
|
||||
@@ -1265,6 +1323,8 @@ int main(int argc, char** argv) {
|
||||
// the next visible ones); "auto" places the K from each GPU's free VRAM once the weights are in (below). Across
|
||||
// GPUs, not yet: KV streaming, images, control vectors, the helper caches (--expert-cache-remote), and lending
|
||||
// cache slots to the prompt path (each stage's prompt path has its own buffers).
|
||||
// --pcie-frac given: that one share is every stage's (the stages' own link probes are skipped), so a split over
|
||||
// a fast and a slow link cannot set the two apart from the command line (#485)
|
||||
const bool pcie_given = o.pcie_frac >= 0.0;
|
||||
std::vector<int64_t> split_at;
|
||||
std::vector<int> split_devs;
|
||||
@@ -1720,15 +1780,14 @@ int main(int argc, char** argv) {
|
||||
// pack's 0.2 was never measured against the link, so it is left alone. `--calibrate` measures it outright.
|
||||
if (o.pcie_frac < 0.0) {
|
||||
const double base = native_pack ? 0.55 : 0.2;
|
||||
const double bw = native_pack ? probe_pcie_h2d_gbps() : -1.0;
|
||||
std::string bursts;
|
||||
const double bw = native_pack ? probe_pcie_h2d_gbps(&bursts) : -1.0;
|
||||
if (!native_pack) {
|
||||
o.pcie_frac = base;
|
||||
} else if (bw > 0.0) {
|
||||
// below ~4 GB/s (an x1 link: ~0.9 GB/s) a missed expert's 1.4 MB takes longer to cross than the CPU
|
||||
// pool takes to compute it, so none of them go over the link
|
||||
o.pcie_frac = bw >= 20.0 ? base : bw < 4.0 ? 0.0 : std::min(base, std::max(0.05, base * (bw / 26.0)));
|
||||
std::fprintf(stderr, "strata generate: PCIe probe: %.1f GB/s host->device -> pcie_frac %.2f (default %.2f)\n",
|
||||
bw, o.pcie_frac, base);
|
||||
o.pcie_frac = pcie_frac_for_gbps(bw, base);
|
||||
std::fprintf(stderr, "strata generate: PCIe probe: %.1f GB/s host->device (best of %s) -> pcie_frac %.2f "
|
||||
"(default %.2f)\n", bw, bursts.c_str(), o.pcie_frac, base);
|
||||
} else {
|
||||
o.pcie_frac = base;
|
||||
std::fprintf(stderr, "strata generate: PCIe probe failed -> pcie_frac default %.2f\n", base);
|
||||
@@ -2229,13 +2288,15 @@ int main(int argc, char** argv) {
|
||||
}
|
||||
strata::kernels::mrope_table_set(st.mrope);
|
||||
}
|
||||
// its own PCIe share of the missed experts (the same rule as CUDA0's above: its link is probed)
|
||||
// its own PCIe share of the missed experts (the same rule as CUDA0's above: its link is probed). A given
|
||||
// --pcie-frac is every stage's share and skips these probes (pcie_given); there is no per-stage setting yet.
|
||||
st.pcie_frac = o.pcie_frac;
|
||||
if (!pcie_given && native_pack) {
|
||||
const double bw = probe_pcie_h2d_gbps();
|
||||
if (bw > 0.0) st.pcie_frac = bw >= 20.0 ? 0.55 : bw < 4.0 ? 0.0 : std::min(0.55, std::max(0.05, 0.55 * (bw / 26.0)));
|
||||
std::fprintf(stderr, "strata generate: layer split: CUDA%d PCIe probe %.1f GB/s -> pcie_frac %.2f\n", st.dev,
|
||||
bw, st.pcie_frac);
|
||||
std::string bursts;
|
||||
const double bw = probe_pcie_h2d_gbps(&bursts);
|
||||
if (bw > 0.0) st.pcie_frac = pcie_frac_for_gbps(bw, 0.55);
|
||||
std::fprintf(stderr, "strata generate: layer split: CUDA%d PCIe probe %.1f GB/s (best of %s) -> pcie_frac "
|
||||
"%.2f\n", st.dev, bw, bursts.c_str(), st.pcie_frac);
|
||||
}
|
||||
size_t fb = 0, tb = 0;
|
||||
cudaMemGetInfo(&fb, &tb);
|
||||
@@ -2258,6 +2319,10 @@ int main(int argc, char** argv) {
|
||||
std::fprintf(stderr, "strata generate: profile %s: %zu ranked pairs, built for %lld slots\n",
|
||||
o.expert_profile.c_str(), profile.size(), (long long) pslots);
|
||||
}
|
||||
// #477: the whole ranking as loaded, the prior of --expert-profile-save's order (a layer split keeps only
|
||||
// CUDA0's pairs in `profile` below). Empty without --expert-profile-save.
|
||||
std::vector<std::pair<int32_t, int32_t>> profile_loaded;
|
||||
if (!o.expert_profile_save.empty()) profile_loaded = profile;
|
||||
// ---- layer split across GPUs: "auto" places the split points by a cost model of one decode window, measured on
|
||||
// the 5080 + 3090 rig (bench/results/2026-09-29-layer-split):
|
||||
// - every layer costs its GPU a time inversely proportional to SMs x clock (0.33 ms on an RTX 5080, 0.50 on a
|
||||
@@ -2553,6 +2618,14 @@ int main(int argc, char** argv) {
|
||||
const char* name = named && p.name[0] ? p.name : "(an unnamed GPU)";
|
||||
#if defined(STRATA_USE_HIP)
|
||||
std::fprintf(stderr, "strata generate: GPU %d: %s (%s)\n", dev, name, named ? p.gcnArchName : "?");
|
||||
#if defined(_WIN32)
|
||||
// #468 #461: which HIP runtime was loaded - the bundled one beside the exe, or an AMD driver's System32 copy
|
||||
if (HMODULE h = GetModuleHandleA("amdhip64_7.dll")) {
|
||||
char path[MAX_PATH] = {};
|
||||
if (GetModuleFileNameA(h, path, MAX_PATH) > 0)
|
||||
std::fprintf(stderr, "strata generate: HIP runtime %s\n", path);
|
||||
}
|
||||
#endif
|
||||
#else
|
||||
std::fprintf(stderr, "strata generate: GPU %d: %s, compute capability %d.%d%s\n", dev, name,
|
||||
strata::cc_major_of(p.major), strata::cc_minor_of(p.minor),
|
||||
@@ -3689,7 +3762,7 @@ int main(int argc, char** argv) {
|
||||
// and each generated token is written to stdout as `T <id>` as soon as its verify window is done, followed by
|
||||
//
|
||||
// DONE <generated> <prompt_tokens> <prompt_ms> <decode_ms> <stop|length|cancel> <drafts accepted>
|
||||
// <drafts offered> <prompt tokens reused>
|
||||
// <drafts offered> <prompt tokens reused> ... <prompt tokens read> (see the DONE line below; #471)
|
||||
//
|
||||
// Before that, `RESUME <n>` (n prompt tokens are not read again), `PP <position> <prompt_tokens> <ms> <tok/s>`
|
||||
// after every prompt chunk, and `REUSED <n>` once the prompt is read. (`ERR <message>` instead when a request
|
||||
@@ -3780,8 +3853,26 @@ int main(int argc, char** argv) {
|
||||
const int64_t k = plan_lend(chunk);
|
||||
if (k > 0) lend_from = xcache.slots() - k;
|
||||
}
|
||||
if (src.pin_cache_complement(xcache, err, o.resident_pin, {}, lend_from, o.resident_headroom,
|
||||
o.resident_budget, &profile)) {
|
||||
bool resident_ok = src.pin_cache_complement(xcache, err, o.resident_pin, {}, lend_from, o.resident_headroom,
|
||||
o.resident_budget, &profile);
|
||||
std::string whole_err;
|
||||
if (!resident_ok && o.resident_soft) {
|
||||
// #467: the whole complement does not fit - keep what does, the hottest by the profile, through the #403
|
||||
// budget path (sized by the RAM alone) instead of none: the misses outside it read the same file bytes
|
||||
// the mmap fallback reads, so the answers are unchanged. Nothing pinned: the old fallback below.
|
||||
whole_err = err;
|
||||
resident_ok = src.pin_cache_complement(xcache, err, o.resident_pin, {}, -1, o.resident_headroom,
|
||||
strata::core::FileExpertSource::kResidentWhatFits, &profile);
|
||||
if (resident_ok)
|
||||
std::fprintf(stderr, "strata generate: WARNING: the whole resident RAM mode does not fit (%s); %.2f "
|
||||
"GiB of the experts the GPU does not hold, the hottest by the expert profile, are "
|
||||
"kept in RAM and the rest are read from the model folder through the OS file "
|
||||
"cache\n",
|
||||
whole_err.c_str(), (double) src.resident_bytes() / 1073741824.0);
|
||||
else
|
||||
err = whole_err + "; " + err;
|
||||
}
|
||||
if (resident_ok) {
|
||||
if (o.adapt_every > 0 && o.adapt_swaps > 0 &&
|
||||
!src.reserve_exchanges(std::min<int64_t>(o.adapt_swaps, 96), err)) {
|
||||
std::fprintf(stderr, "strata generate: CPU expert residency: %s\n", err.c_str());
|
||||
@@ -3887,24 +3978,62 @@ int main(int argc, char** argv) {
|
||||
// 128-slot floor. The percentage cap is an AUTO-chunk rule and only the auto scan applies it - an
|
||||
// explicit --prefill is the operator's number, and a loan of it only has to fit. With one participant
|
||||
// (no split) this reduces to plan_lend exactly, so the single-GPU loan is unchanged from main.
|
||||
auto fits = [&](int64_t c, bool cap) -> bool {
|
||||
for (const PfPart& p : pf_parts) {
|
||||
const int64_t k = part_slots(p, c);
|
||||
if (k <= 0 || k + 128 > p.cache->slots()) return false;
|
||||
if (cap && k * 100 > kAutoLendPct * p.cache->slots()) return false;
|
||||
}
|
||||
auto fits_one = [&](const PfPart& p, int64_t c, bool cap) -> bool {
|
||||
const int64_t k = part_slots(p, c);
|
||||
if (k <= 0 || k + 128 > p.cache->slots()) return false;
|
||||
return !(cap && k * 100 > kAutoLendPct * p.cache->slots());
|
||||
};
|
||||
// `only`: CUDA0's cache alone (#448: what one GPU would choose, for the log below); null: every one
|
||||
auto fits = [&](int64_t c, bool cap, const PfPart* only = nullptr) -> bool {
|
||||
if (only != nullptr) return fits_one(*only, c, cap);
|
||||
for (const PfPart& p : pf_parts)
|
||||
if (!fits_one(p, c, cap)) return false;
|
||||
return true;
|
||||
};
|
||||
static constexpr int64_t kAutoChunks[] = {32768, 16384, 8192, 6144, 4096, 3072, 2048, 1024, 512, 256};
|
||||
int64_t chunk = 0;
|
||||
if (o.prefill_auto) {
|
||||
for (const int64_t c : kAutoChunks) {
|
||||
if (c > 8192 && (c > o.prefill_auto_max || c > o.max_context)) continue; // #282, as plan_lend
|
||||
if (fits(c, true)) { chunk = c; break; }
|
||||
auto pick = [&](const PfPart* only) -> int64_t {
|
||||
if (o.prefill_auto) {
|
||||
for (const int64_t c : kAutoChunks) {
|
||||
if (c > 8192 && (c > o.prefill_auto_max || c > o.max_context)) continue; // #282, as plan_lend
|
||||
if (fits(c, true, only)) return c;
|
||||
}
|
||||
} else {
|
||||
for (int64_t c = o.prefill_chunk; c >= 256; c /= 2)
|
||||
if (fits(c, false, only)) return c;
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
const int64_t chunk = pick(nullptr);
|
||||
// #448: a small card in a layer split caps every stage's chunk (an RTX 3080's 512-slot cache held a
|
||||
// 32 GB card's split to 512 tokens: prompts 6.2x slower, decode the same). Named when it bites, so the
|
||||
// regression is one log line: each stage that cannot fund the chunk CUDA0 alone would read in.
|
||||
if (pf_parts.size() > 1) {
|
||||
const int64_t alone = pick(&pf_parts[0]);
|
||||
if (alone > chunk) {
|
||||
for (size_t i = 1; i < pf_parts.size(); ++i) {
|
||||
const PfPart& p = pf_parts[i];
|
||||
if (fits_one(p, alone, o.prefill_auto)) continue;
|
||||
const int dev = p.dev < 0 ? 0 : p.dev;
|
||||
cudaDeviceProp prop{};
|
||||
if (cudaGetDeviceProperties(&prop, dev) != cudaSuccess) {
|
||||
(void) cudaGetLastError();
|
||||
prop.name[0] = 0;
|
||||
}
|
||||
const std::string pct =
|
||||
o.prefill_auto ? ", and lend at most " + std::to_string(kAutoLendPct) + "%" : "";
|
||||
std::fprintf(stderr, "strata serve: WARNING: prompt chunk %lld tokens, not %lld: CUDA%d (%s) "
|
||||
"has %lld expert-cache slots, and a %lld-token chunk borrows %lld of them "
|
||||
"(it must keep 128%s) - prompts read slower than on CUDA0 alone (#448)\n",
|
||||
(long long) chunk, (long long) alone, dev, prop.name,
|
||||
(long long) p.cache->slots(), (long long) alone,
|
||||
(long long) part_slots(p, alone), pct.c_str());
|
||||
// the helper tiers start at CUDA1 without a split (and are enabled in order)
|
||||
std::fprintf(stderr, "strata serve: a card this small can serve as a helper expert cache "
|
||||
"instead of a split stage: without --layer-split, with %s "
|
||||
"(docs/SECOND_GPU.md)\n",
|
||||
dev == 1 ? "--expert-cache-device1 N" : "--expert-cache-device1..3 N, in order");
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for (int64_t c = o.prefill_chunk; c >= 256; c /= 2)
|
||||
if (fits(c, false)) { chunk = c; break; }
|
||||
}
|
||||
if (chunk > 0) {
|
||||
if (o.prefill_auto)
|
||||
@@ -4309,6 +4438,9 @@ int main(int argc, char** argv) {
|
||||
return true;
|
||||
};
|
||||
int64_t pp_total = 0, pp_from = 0, pp_next_check = 0;
|
||||
// #471: the position the prompt pass has read up to (a chunk's or a window's end): what a request cancelled
|
||||
// mid-read reports as read, instead of the whole prompt
|
||||
int64_t pp_reached = 0;
|
||||
Clock::time_point pp_t0 = Clock::now();
|
||||
auto imgs_below = [&](const std::vector<ImgKey>& all, int64_t L) {
|
||||
std::vector<ImgKey> v;
|
||||
@@ -4372,6 +4504,7 @@ int main(int argc, char** argv) {
|
||||
batched ? "batched" : "token", mtp.kv_state().kv_mode, (long long) T);
|
||||
// progress for the server window: PP <position reached> <prompt tokens> <ms> <fresh tokens/s>
|
||||
const int64_t done = p0 + T;
|
||||
pp_reached = done;
|
||||
const double ms = std::chrono::duration<double, std::milli>(Clock::now() - pp_t0).count();
|
||||
std::printf("PP %lld %lld %.0f %.1f\n", (long long) done, (long long) pp_total, ms,
|
||||
ms > 0.0 ? 1000.0 * (double) (done - pp_from) / ms : 0.0);
|
||||
@@ -4422,6 +4555,17 @@ int main(int argc, char** argv) {
|
||||
drive.d.plan = ver.plan_sink();
|
||||
drive.d.pcie_num = std::max(0, std::min(256, (int) (o.pcie_frac * 256.0 + 0.5)));
|
||||
if (o.adapt_every > 0 && o.adapt_swaps > 0) drive.d.usage.assign((size_t) (g.n_layers * g.n_expert), 0.0f);
|
||||
// #477 --expert-profile-save: what the adaptive tier learned, kept across restarts (opt-in; off: `heat` stays
|
||||
// empty and nothing below runs). It needs the adaptive tier's counts and the residency table.
|
||||
std::vector<double> heat;
|
||||
if (!o.expert_profile_save.empty()) {
|
||||
if (drive.d.usage.empty() || host_res.empty())
|
||||
std::fprintf(stderr, "strata serve: --expert-profile-save needs the adaptive tier (--adapt-every and "
|
||||
"--adapt-swaps above 0) and --expert-profile: nothing will be saved\n");
|
||||
else
|
||||
heat.assign(drive.d.usage.size(), 0.0);
|
||||
}
|
||||
Clock::time_point profile_saved_at = Clock::now();
|
||||
cudaStream_t adapt_stream = nullptr;
|
||||
if (cudaStreamCreateWithFlags(&adapt_stream, cudaStreamNonBlocking) != cudaSuccess) {
|
||||
std::fprintf(stderr, "strata serve: cannot create the refill stream\n");
|
||||
@@ -4508,9 +4652,29 @@ int main(int argc, char** argv) {
|
||||
const strata::core::OnDevice on(st->dev);
|
||||
cudaEventRecord(st->adapt_ev, st->adapt_stream);
|
||||
}
|
||||
// #477: the routing counted since the start (each count adds up to 1 / (1 - 0.7) over its decays: the
|
||||
// sum is proportional to the routing itself) - only with --expert-profile-save, else `heat` is empty
|
||||
for (size_t i = 0; i < heat.size(); ++i) heat[i] += (double) drive.d.usage[i];
|
||||
for (float& v : drive.d.usage) v *= 0.7f;
|
||||
return true;
|
||||
};
|
||||
// #477: write the learned profile (between requests and at QUIT: a prompt's lent slots are back by then).
|
||||
// A swap still in flight counts as done - its expert is resident once the copy lands. `why`: for the log.
|
||||
auto save_profile = [&](const char* why) {
|
||||
if (heat.empty()) return;
|
||||
std::vector<uint8_t> resident(host_res.size(), 0);
|
||||
for (size_t i = 0; i < host_res.size(); ++i) resident[i] = host_res[i] >= 0;
|
||||
for (const auto& p : pending) resident[(size_t) p.first] = 1;
|
||||
std::string e;
|
||||
const auto ranked = strata::core::rank_learned_profile(g.n_layers, g.n_expert, resident, heat,
|
||||
profile_loaded);
|
||||
if (strata::core::write_expert_profile(o.expert_profile_save, g.n_layers, g.n_expert, ranked, e))
|
||||
std::fprintf(stderr, "strata serve: expert profile saved to %s (%s)\n", o.expert_profile_save.c_str(),
|
||||
why);
|
||||
else
|
||||
std::fprintf(stderr, "strata serve: the expert profile was not saved: %s\n", e.c_str());
|
||||
profile_saved_at = Clock::now();
|
||||
};
|
||||
// stdin is read on its own thread, so a STOP line reaches a request that is still running (the client went
|
||||
// away, or pressed Esc): the flag is checked between prompt chunks and between verify windows.
|
||||
std::atomic<bool> stop_req{false};
|
||||
@@ -4673,6 +4837,10 @@ int main(int argc, char** argv) {
|
||||
std::vector<float> img_rows;
|
||||
std::vector<const float*> row_ptr;
|
||||
while (next_line(line)) {
|
||||
// #477: every --expert-profile-save-every minutes, before the next request (at QUIT: after the loop)
|
||||
if (!heat.empty() && line != "QUIT" && o.expert_profile_save_min > 0 &&
|
||||
Clock::now() - profile_saved_at >= std::chrono::duration<double>(o.expert_profile_save_min * 60.0))
|
||||
save_profile("periodic");
|
||||
if (line == "QUIT") break;
|
||||
// the watchdog watches a request from here until this iteration ends, whichever way it ends
|
||||
struct BusyScope {
|
||||
@@ -4999,6 +5167,7 @@ int main(int argc, char** argv) {
|
||||
conversations.limit_reuse(read_from);
|
||||
pp_total = n;
|
||||
pp_from = read_from;
|
||||
pp_reached = read_from;
|
||||
pp_t0 = r0;
|
||||
pp_next_check = reread_to > 0 ? INT64_MAX : resume + o.prompt_cache_every;
|
||||
{
|
||||
@@ -5064,6 +5233,7 @@ int main(int argc, char** argv) {
|
||||
return false;
|
||||
if (!ver.commit(T, e) || !mtp.prefill(ver.final_R_all(), nxt.data(), T, q, e)) return false;
|
||||
q += T;
|
||||
pp_reached = q; // #471
|
||||
}
|
||||
// the batched prompt path (other streams), checkpoints and snapshots may follow: the last commit first
|
||||
if (!ver.wait_commit(e)) return false;
|
||||
@@ -5564,20 +5734,30 @@ int main(int argc, char** argv) {
|
||||
}
|
||||
const int64_t req_hits = drive.d.cache_hits - decode_hits0;
|
||||
const int64_t req_look = (drive.d.cache_hits + drive.d.cache_admitted + drive.d.cache_refused) - decode_look0;
|
||||
// #471: the prompt tokens this request read - all the fresh ones, or as far as the prompt pass got when a
|
||||
// cancel stopped it part-way (a cancelled request used to be logged and counted as having read them all)
|
||||
const int64_t fresh = n - resume;
|
||||
const int64_t read_n = cancelled ? std::clamp<int64_t>(pp_reached - resume, 0, fresh) : fresh;
|
||||
// DONE <generated> <prompt> <prompt ms> <decode ms> <finish> <drafts accepted> <drafts offered> <reused> [hits] [lookups]
|
||||
// [RAM blobs] [file blobs] [file MB] (CS-T tiers; appended, so an older server reads the rest)
|
||||
std::printf("DONE %lld %lld %.1f %.1f %s %lld %lld %lld %lld %lld %lld %lld %.1f\n", (long long) produced_n,
|
||||
// [prompt tokens read] (#471: fewer than <prompt> - <reused> when a cancel stopped the read)
|
||||
std::printf("DONE %lld %lld %.1f %.1f %s %lld %lld %lld %lld %lld %lld %lld %.1f %lld\n", (long long) produced_n,
|
||||
(long long) n, prompt_ms, decode_ms, finish, (long long) draft_accepted, (long long) draft_offered,
|
||||
(long long) resume, (long long) req_hits, (long long) req_look,
|
||||
(long long) (src.ram_reads() - ram0), (long long) (src.file_reads() - files0),
|
||||
(double) (src.file_read_bytes() - file_bytes0) / 1e6);
|
||||
(double) (src.file_read_bytes() - file_bytes0) / 1e6, (long long) read_n);
|
||||
std::fflush(stdout);
|
||||
if (drive.routing != nullptr) std::fflush(drive.routing); // the routing trace survives a crash and is watchable mid-session
|
||||
const int64_t fresh = n - resume;
|
||||
std::fprintf(stderr, "strata serve: prompt %lld tokens = %lld reused + %lld read in %.0f ms (%.1f tok/s), "
|
||||
// "12288 of 98179" when cancelled mid-read (#471), the rate from what was read
|
||||
char read_txt[64];
|
||||
if (cancelled)
|
||||
std::snprintf(read_txt, sizeof(read_txt), "%lld of %lld", (long long) read_n, (long long) fresh);
|
||||
else
|
||||
std::snprintf(read_txt, sizeof(read_txt), "%lld", (long long) fresh);
|
||||
std::fprintf(stderr, "strata serve: prompt %lld tokens = %lld reused + %s read in %.0f ms (%.1f tok/s), "
|
||||
"%lld generated in %.0f ms (%.1f tok/s), drafts accepted %lld of %lld, %zu checkpoints%s\n",
|
||||
(long long) n, (long long) resume, (long long) fresh, prompt_ms,
|
||||
prompt_ms > 0 ? 1000.0 * fresh / prompt_ms : 0.0, (long long) produced_n, decode_ms,
|
||||
(long long) n, (long long) resume, read_txt, prompt_ms,
|
||||
prompt_ms > 0 ? 1000.0 * read_n / prompt_ms : 0.0, (long long) produced_n, decode_ms,
|
||||
decode_ms > 0 ? 1000.0 * produced_n / decode_ms : 0.0, (long long) draft_accepted,
|
||||
(long long) draft_offered, checks.size(), cancelled ? " (cancelled)" : "");
|
||||
// the VRAM share of the experts the pool looked up while decoding; experts it sent over PCIe for the GPU
|
||||
@@ -5639,6 +5819,7 @@ int main(int argc, char** argv) {
|
||||
remote_experts[(size_t) r].ms_begin() - begin_before[(size_t) r],
|
||||
remote_experts[(size_t) r].ms_wait() - wait_before[(size_t) r]);
|
||||
}
|
||||
save_profile("exit"); // #477: QUIT, or the server closed stdin
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
// #477: --expert-profile-save's profile - rank_learned_profile's order (resident first, then the counted routing, then
|
||||
// the profile the engine started from) and write_expert_profile's file, byte for byte tools/make_profile.py's format,
|
||||
// read back by read_expert_profile. CPU only: no device is touched.
|
||||
#include "strata/core/expert_cache.hpp"
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <filesystem>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
int fails = 0;
|
||||
void check(bool ok, const char* what) {
|
||||
if (!ok) {
|
||||
std::fprintf(stderr, "FAIL: %s\n", what);
|
||||
++fails;
|
||||
}
|
||||
}
|
||||
} // namespace
|
||||
|
||||
int main() {
|
||||
using Pair = std::pair<int32_t, int32_t>;
|
||||
const int64_t L = 3, E = 4; // 12 pairs, index = layer * 4 + expert
|
||||
std::vector<uint8_t> resident(12, 0);
|
||||
resident[5] = 1; // (1, 1)
|
||||
resident[2] = 1; // (0, 2)
|
||||
std::vector<double> heat(12, 0.0);
|
||||
heat[2] = 1.0;
|
||||
heat[5] = 7.0;
|
||||
heat[11] = 9.0; // (2, 3): the hottest, but not resident
|
||||
heat[0] = 3.0;
|
||||
const std::vector<Pair> prior = {{2, 0}, {1, 3}, {0, 1}}; // a profile that ranks three pairs
|
||||
const std::vector<Pair> r = strata::core::rank_learned_profile(L, E, resident, heat, prior);
|
||||
check(r.size() == 12, "every pair is ranked");
|
||||
const std::vector<Pair> head = {{1, 1}, {0, 2}, {2, 3}, {0, 0}, {2, 0}, {1, 3}, {0, 1}, {0, 3}};
|
||||
for (size_t i = 0; i < head.size() && i < r.size(); ++i)
|
||||
check(r[i] == head[i], "resident by heat, then the rest by heat, then the prior, then the index");
|
||||
// no counts at all: the resident ones, then the prior's order
|
||||
const std::vector<Pair> r0 = strata::core::rank_learned_profile(L, E, resident, {}, prior);
|
||||
check(r0.size() == 12 && r0[0] == Pair(0, 2) && r0[1] == Pair(1, 1) && r0[2] == Pair(2, 0) &&
|
||||
r0[3] == Pair(1, 3) && r0[4] == Pair(0, 1) && r0[5] == Pair(0, 0), "without heat: resident, then prior");
|
||||
|
||||
const std::filesystem::path dir = std::filesystem::temp_directory_path() / "strata_profile_save_test";
|
||||
std::filesystem::create_directories(dir);
|
||||
const std::string path = (dir / "learned.bin").string();
|
||||
std::string err;
|
||||
check(strata::core::write_expert_profile(path, L, E, r, err), "written");
|
||||
check(!std::filesystem::exists(path + ".tmp"), "no temporary file left");
|
||||
// the bytes make_profile.write_profile would write for the same ranking
|
||||
std::vector<uint8_t> want(4);
|
||||
std::memcpy(want.data(), "STRP", 4);
|
||||
auto u32 = [&](uint32_t v) { for (int b = 0; b < 4; ++b) want.push_back((uint8_t) (v >> (8 * b))); };
|
||||
auto u16 = [&](uint32_t v) { want.push_back((uint8_t) v); want.push_back((uint8_t) (v >> 8)); };
|
||||
for (uint32_t v : {1u, 3u, 4u, 12u, 12u}) u32(v);
|
||||
for (const Pair& p : r) { u16((uint32_t) p.first); u16((uint32_t) p.second); }
|
||||
std::vector<int32_t> table(12, -1);
|
||||
for (size_t i = 0; i < r.size(); ++i) table[(size_t) (r[i].first * E + r[i].second)] = (int32_t) i;
|
||||
for (int32_t v : table) u32((uint32_t) v);
|
||||
std::vector<uint8_t> got(want.size() + 1);
|
||||
if (std::FILE* f = std::fopen(path.c_str(), "rb")) {
|
||||
got.resize(std::fread(got.data(), 1, got.size(), f));
|
||||
std::fclose(f);
|
||||
}
|
||||
check(got == want, "the file is make_profile.py's format, byte for byte");
|
||||
std::vector<Pair> back;
|
||||
int64_t slots = 0;
|
||||
check(strata::core::read_expert_profile(path, L, E, back, slots, err) && back == r && slots == 12,
|
||||
"read back by the loader");
|
||||
// a second save replaces the first (the rename over an existing file)
|
||||
check(strata::core::write_expert_profile(path, L, E, r0, err), "written again");
|
||||
check(strata::core::read_expert_profile(path, L, E, back, slots, err) && back == r0, "replaced");
|
||||
check(!strata::core::write_expert_profile(path, L, E, {{3, 0}}, err), "a pair out of range is refused");
|
||||
check(strata::core::read_expert_profile(path, L, E, back, slots, err) && back == r0,
|
||||
"a refused save keeps the file");
|
||||
std::filesystem::remove_all(dir);
|
||||
if (fails == 0) std::puts("expert_profile_save_test: OK");
|
||||
return fails == 0 ? 0 : 1;
|
||||
}
|
||||
@@ -0,0 +1,633 @@
|
||||
// prefill_fused_iq_test - #136: the fused int8 expert kernels for the native packs (moe_fused_iq.hpp, STRATA_PF_FUSED=1)
|
||||
// on random native expert blobs ([gate rows | up rows | down rows], GGUF blocks) and random routing, per format pair,
|
||||
// against
|
||||
// - a double-precision reference: the weights dequantized by ggml's own to_float, times the FP32 activations
|
||||
// (gate/up, SwiGLU, down), and
|
||||
// - the MMQ path in prefill.cpp's sequence (gather_native into 16-expert groups, q8_1 of the slots' activations,
|
||||
// gate/up, SwiGLU, q8_1 of H, down).
|
||||
// Both round the activations and H to int8 per 32 values, differently, so this is not a bitwise test: the fused
|
||||
// path's error against the reference has to be comparable to MMQ's own (at most 1.5x its RMS and 2x its worst row).
|
||||
// The pairs cover every format the kernels take: gate/up IQ2_XXS IQ2_XS IQ2_S IQ3_XXS IQ3_S IQ4_XS, down Q2_0 IQ4_NL.
|
||||
// Part 1 (reference), per pair: 256 tokens over 64 experts, skewed routing - experts with 0 rows, with one, with
|
||||
// several 64-row tiles - an all-zero token, per-expert blobs at unrelated (2-byte aligned) addresses, three launches of
|
||||
// different expert ranges. Part 2 (timing): one layer at the engine's chunks (2048, 3584, 8192 tokens; 512 experts,
|
||||
// top 10) on both paths and on MMQ's products alone (no gathers), for the IQ2_XS, IQ3_XXS and IQ3_S packs' most
|
||||
// common layers, with the grouping checked and a sample of rows against the reference. --no-ref, --no-timing,
|
||||
// --only=NAME, --chunks=A,B. Exit 77 without a CUDA device of sm_80 or newer.
|
||||
#include "strata/prefill/moe_fused_iq.hpp"
|
||||
#include "strata/prefill/moe_mmq.hpp"
|
||||
|
||||
#include "ggml.h"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <random>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
namespace mmq = strata::prefill::mmq;
|
||||
namespace fused = strata::prefill::fused;
|
||||
|
||||
constexpr int N = 2560, FF = 640, K = 10, GROUP = 16;
|
||||
|
||||
void ck(cudaError_t e, const char* what) {
|
||||
if (e != cudaSuccess) throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(e));
|
||||
}
|
||||
|
||||
struct Dev {
|
||||
void* p = nullptr;
|
||||
explicit Dev(size_t n) { ck(cudaMalloc(&p, n), "cudaMalloc"); }
|
||||
~Dev() { cudaFree(p); }
|
||||
Dev(const Dev&) = delete;
|
||||
Dev& operator=(const Dev&) = delete;
|
||||
template <typename T> T* as() const { return (T*) p; }
|
||||
};
|
||||
|
||||
struct Pair {
|
||||
ggml_type gu, d;
|
||||
const char* name;
|
||||
};
|
||||
|
||||
struct Geo {
|
||||
size_t gu_row, d_row, up_off, down_off, bytes;
|
||||
explicit Geo(const Pair& p) {
|
||||
gu_row = ggml_row_size(p.gu, N);
|
||||
d_row = ggml_row_size(p.d, FF);
|
||||
up_off = gu_row * FF;
|
||||
down_off = 2 * up_off;
|
||||
bytes = down_off + d_row * N;
|
||||
}
|
||||
fused::NativeGeom native(const Pair& p) const {
|
||||
fused::NativeGeom g;
|
||||
g.gu_type = (int) p.gu; g.d_type = (int) p.d; g.gu_row = gu_row; g.d_row = d_row; g.up_off = up_off;
|
||||
g.down_off = down_off;
|
||||
return g;
|
||||
}
|
||||
};
|
||||
|
||||
// random blocks with a sane fp16 scale `d` at each block's start (every format here begins with it): scales chosen so
|
||||
// that the weights' typical size is that of a 2-3 bit expert (a few hundredths)
|
||||
void fill_matrix(uint8_t* m, ggml_type t, int64_t rows, int64_t cols, std::mt19937& rng) {
|
||||
const size_t bs = ggml_type_size(t), nb = (size_t) rows * (size_t) (cols / ggml_blck_size(t));
|
||||
std::uniform_int_distribution<int> byte(0, 255);
|
||||
const float lo = t == GGML_TYPE_IQ3_S ? 0.0004f : t == GGML_TYPE_IQ4_XS ? 0.00005f
|
||||
: t == GGML_TYPE_IQ4_NL ? 0.0002f : t == GGML_TYPE_Q2_0 ? 0.004f : 0.0008f;
|
||||
std::uniform_real_distribution<float> sc(lo, 6.0f * lo);
|
||||
for (size_t i = 0; i < nb * bs; ++i) m[i] = (uint8_t) byte(rng);
|
||||
for (size_t b = 0; b < nb; ++b) {
|
||||
const ggml_fp16_t h = ggml_fp32_to_fp16(sc(rng));
|
||||
std::memcpy(m + b * bs, &h, 2);
|
||||
}
|
||||
}
|
||||
std::vector<uint8_t> make_blob(const Pair& p, const Geo& g, std::mt19937& rng) {
|
||||
std::vector<uint8_t> b(g.bytes);
|
||||
fill_matrix(b.data(), p.gu, 2 * FF, N, rng);
|
||||
fill_matrix(b.data() + g.down_off, p.d, N, FF, rng);
|
||||
return b;
|
||||
}
|
||||
|
||||
// ggml's dequantized weights: gate [640][2560], up [640][2560], down [2560][640]
|
||||
void dequant(const Pair& p, const Geo& g, const std::vector<uint8_t>& b, std::vector<float>& gu, std::vector<float>& dn) {
|
||||
const auto* tg = ggml_get_type_traits(p.gu);
|
||||
const auto* td = ggml_get_type_traits(p.d);
|
||||
gu.resize((size_t) 2 * FF * N);
|
||||
dn.resize((size_t) N * FF);
|
||||
for (int r = 0; r < 2 * FF; ++r) tg->to_float(b.data() + (size_t) r * g.gu_row, gu.data() + (size_t) r * N, N);
|
||||
for (int r = 0; r < N; ++r) td->to_float(b.data() + g.down_off + (size_t) r * g.d_row, dn.data() + (size_t) r * FF, FF);
|
||||
}
|
||||
|
||||
// activations: N(0, 1) with a few large values (the hidden state has outliers); token `zero` all zero
|
||||
std::vector<float> make_x(int T, std::mt19937& rng, int zero) {
|
||||
std::vector<float> x((size_t) T * N);
|
||||
std::normal_distribution<float> nd(0.0f, 1.0f);
|
||||
std::uniform_int_distribution<int> pick(0, 199);
|
||||
for (float& v : x) v = nd(rng) * (pick(rng) == 0 ? 12.0f : 1.0f);
|
||||
if (zero >= 0) std::fill(x.begin() + (size_t) zero * N, x.begin() + (size_t) (zero + 1) * N, 0.0f);
|
||||
return x;
|
||||
}
|
||||
|
||||
// K distinct experts per token; `hot` > 0 skews the first choices towards the lowest ids
|
||||
std::vector<int32_t> make_ids(int T, int E, std::mt19937& rng, int hot, int unused) {
|
||||
std::vector<int32_t> ids((size_t) T * K);
|
||||
std::uniform_int_distribution<int> any(0, E - 1 - unused);
|
||||
for (int t = 0; t < T; ++t)
|
||||
for (int k = 0; k < K; ++k) {
|
||||
int e;
|
||||
bool dup;
|
||||
do {
|
||||
e = (hot > 0 && k < 3 && (int) (rng() % 4) != 0) ? (int) (rng() % (unsigned) hot) : any(rng);
|
||||
dup = false;
|
||||
for (int j = 0; j < k; ++j) dup |= ids[(size_t) t * K + j] == e;
|
||||
} while (dup);
|
||||
ids[(size_t) t * K + k] = e;
|
||||
}
|
||||
return ids;
|
||||
}
|
||||
|
||||
struct Routing {
|
||||
std::vector<int32_t> cnt, off, src; // per expert; sorted row -> token (rows in pair order within an expert)
|
||||
std::vector<int32_t> row_of; // pair -> sorted row
|
||||
};
|
||||
Routing sort_rows(const std::vector<int32_t>& ids, int E) {
|
||||
Routing r;
|
||||
r.cnt.assign((size_t) E, 0);
|
||||
for (int32_t e : ids) ++r.cnt[(size_t) e];
|
||||
r.off.assign((size_t) E + 1, 0);
|
||||
for (int e = 0; e < E; ++e) r.off[(size_t) e + 1] = r.off[(size_t) e] + r.cnt[(size_t) e];
|
||||
std::vector<int32_t> fill(r.off.begin(), r.off.end() - 1);
|
||||
r.src.resize(ids.size());
|
||||
r.row_of.resize(ids.size());
|
||||
for (size_t i = 0; i < ids.size(); ++i) {
|
||||
const int32_t p = fill[(size_t) ids[i]]++;
|
||||
r.row_of[i] = p;
|
||||
r.src[(size_t) p] = (int32_t) (i / K);
|
||||
}
|
||||
return r;
|
||||
}
|
||||
|
||||
// MMQ, as prefill.cpp runs a native layer: per group of 16 routed experts a gather, gate/up, SwiGLU (gate rows, then
|
||||
// up rows), q8_1 of H, down. Returns Dm [rows][N] in sorted-row order.
|
||||
struct MmqBufs {
|
||||
Dev xq, gu, h, hq, dm, ident, bounds, grp_gu, grp_d;
|
||||
MmqBufs(const Pair& p, int64_t rows, int E)
|
||||
: xq(mmq::q8_bytes(rows, N)), gu((size_t) rows * 1280 * 4), h((size_t) rows * FF * 4),
|
||||
hq(mmq::q8_bytes(rows, FF)), dm((size_t) rows * N * 4), ident((size_t) rows * 4),
|
||||
bounds((size_t) (2 * (E + E / GROUP + 2)) * 4), grp_gu(GROUP * mmq::matrix_bytes(p.gu, 1280, N) + 4096),
|
||||
grp_d(GROUP * mmq::matrix_bytes(p.d, N, FF) + 4096) {}
|
||||
};
|
||||
// `direct`: no gathers - the products read the blobs where they are (they must lie at a fixed stride, every expert
|
||||
// with rows: the timing part's layout), the cost of MMQ's own kernels alone
|
||||
void run_mmq(const Pair& p, const Geo& geo, mmq::Context& ctx, MmqBufs& b, const Routing& r, const float* x_dev,
|
||||
const int32_t* src_dev, const std::vector<const uint8_t*>& blob, int64_t rows, cudaStream_t s,
|
||||
bool direct = false, size_t stride = 0) {
|
||||
const size_t gub = mmq::matrix_bytes(p.gu, 1280, N), db = mmq::matrix_bytes(p.d, N, FF);
|
||||
const int E = (int) r.cnt.size();
|
||||
std::vector<int32_t> order;
|
||||
for (int e = 0; e < E; ++e) if (r.cnt[(size_t) e] > 0) order.push_back(e);
|
||||
const size_t n = order.size(), ng = (n + GROUP - 1) / GROUP;
|
||||
std::vector<int32_t> bh(n + 1 + ng * (GROUP + 1));
|
||||
for (size_t j = 0; j < n; ++j) bh[j] = r.off[(size_t) order[j]];
|
||||
bh[n] = (int32_t) rows;
|
||||
for (size_t g = 0; g < ng; ++g)
|
||||
for (size_t i = 0; i <= GROUP; ++i)
|
||||
bh[n + 1 + g * (GROUP + 1) + i] = bh[std::min(n, g * GROUP + i)] - bh[g * GROUP];
|
||||
ck(cudaMemcpyAsync(b.bounds.p, bh.data(), bh.size() * 4, cudaMemcpyHostToDevice, s), "bounds");
|
||||
mmq::iota(b.ident.as<int32_t>(), rows, s);
|
||||
mmq::quantize(x_dev, src_dev, b.xq.p, (int) p.gu, N, N, rows, s);
|
||||
for (size_t j = 0; j < n; ++j) {
|
||||
const size_t q = j % GROUP;
|
||||
const uint8_t* bl = blob[(size_t) order[j]];
|
||||
if (!direct)
|
||||
mmq::gather_native(bl, bl + geo.up_off, gub / 2, bl + geo.down_off, db, b.grp_gu.as<uint8_t>() + q * gub,
|
||||
b.grp_d.as<uint8_t>() + q * db, s);
|
||||
if (q + 1 < GROUP && j + 1 < n) continue;
|
||||
const size_t j0 = j - q, g = j0 / GROUP;
|
||||
const int ngx = (int) (q + 1);
|
||||
const int64_t r0 = bh[j0], nr = bh[j + 1] - r0;
|
||||
int64_t maxr = 0;
|
||||
for (size_t i = j0; i <= j; ++i) maxr = std::max<int64_t>(maxr, r.cnt[(size_t) order[i]]);
|
||||
if (!direct) {
|
||||
ck(cudaMemsetAsync(b.grp_gu.as<uint8_t>() + (size_t) ngx * gub, 0, 4096, s), "tail");
|
||||
ck(cudaMemsetAsync(b.grp_d.as<uint8_t>() + (size_t) ngx * db, 0, 4096, s), "tail");
|
||||
}
|
||||
const uint8_t* b0 = blob[(size_t) order[j0]];
|
||||
mmq::Product gu;
|
||||
gu.w = direct ? (const void*) b0 : b.grp_gu.p; gu.type = (int) p.gu; gu.w_rows = 1280; gu.w_cols = N;
|
||||
gu.expert_bytes = direct ? stride : gub;
|
||||
gu.n = ngx; gu.xq = b.xq.p; gu.bounds = b.bounds.as<int32_t>() + j0; gu.ids = b.ident.as<int32_t>();
|
||||
gu.total_rows = rows; gu.max_rows = maxr; gu.dst = b.gu.as<float>(); gu.ld_dst = 1280;
|
||||
ctx.run(gu, s);
|
||||
mmq::swiglu(b.gu.as<float>() + r0 * 1280, b.h.as<float>() + r0 * FF, nr, FF, false, s);
|
||||
mmq::quantize(b.h.as<float>() + r0 * FF, nullptr, b.hq.p, (int) p.d, FF, FF, nr, s);
|
||||
mmq::Product dn;
|
||||
dn.w = direct ? (const void*) (b0 + geo.down_off) : b.grp_d.p; dn.type = (int) p.d; dn.w_rows = N; dn.w_cols = FF;
|
||||
dn.expert_bytes = direct ? stride : db;
|
||||
dn.n = ngx; dn.xq = b.hq.p; dn.bounds = b.bounds.as<int32_t>() + n + 1 + g * (GROUP + 1);
|
||||
dn.ids = b.ident.as<int32_t>(); dn.total_rows = nr; dn.max_rows = maxr; dn.dst = b.dm.as<float>() + r0 * N;
|
||||
dn.ld_dst = N;
|
||||
ctx.run(dn, s);
|
||||
}
|
||||
}
|
||||
|
||||
// the fused path: quantize, group on the device, then launches over expert ranges `cuts` (e.g. {0, 20, 64})
|
||||
struct FusedBufs {
|
||||
Dev xa, ha, dm, scratch, slot, src, ids;
|
||||
FusedBufs(int T, int64_t rows, int E)
|
||||
: xa(fused::act_bytes(T, N)), ha(fused::act_bytes(rows, FF)), dm((size_t) rows * N * 4),
|
||||
scratch(fused::group_bytes(rows, E)), slot((size_t) rows * 4), src((size_t) rows * 4),
|
||||
ids((size_t) rows * 4) {}
|
||||
};
|
||||
void run_fused(const Pair& p, const Geo& geo, FusedBufs& b, const float* x_dev, int T, int E,
|
||||
const std::vector<int>& cuts, const std::vector<const uint8_t*>& blob, cudaStream_t s) {
|
||||
const int64_t rows = (int64_t) T * K;
|
||||
const fused::NativeGeom ng = geo.native(p);
|
||||
fused::quantize_act_native(x_dev, T, N, b.xa.p, s);
|
||||
fused::group(b.ids.as<int32_t>(), rows, K, E, b.scratch.p, b.slot.as<int32_t>(), b.src.as<int32_t>(), s);
|
||||
for (size_t c = 0; c + 1 < cuts.size(); ++c) {
|
||||
for (int e0 = cuts[c]; e0 < cuts[c + 1]; e0 += fused::kMaxBatch) {
|
||||
fused::Batch bt;
|
||||
bt.e0 = e0;
|
||||
bt.e1 = std::min(cuts[c + 1], e0 + fused::kMaxBatch);
|
||||
for (int e = bt.e0; e < bt.e1; ++e) bt.blob[e - bt.e0] = blob[(size_t) e];
|
||||
fused::experts_native(bt, ng, E, rows, b.scratch.p, b.xa.p, b.src.as<int32_t>(), b.ha.p, b.dm.as<float>(), s);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct Err {
|
||||
double rms = 0, worst = 0; // RMS error / RMS of the reference; the worst row's max |error| / its max |ref|
|
||||
};
|
||||
// y_a[row_a(i)] against y_b[row_b(i)] for every pair i (rows of N values)
|
||||
template <typename RA, typename RB>
|
||||
Err compare(const std::vector<float>& a, RA row_a, const std::vector<float>& b, RB row_b, size_t pairs) {
|
||||
double e2 = 0, r2 = 0;
|
||||
Err out;
|
||||
for (size_t i = 0; i < pairs; ++i) {
|
||||
const float* ya = a.data() + (size_t) row_a(i) * N;
|
||||
const float* yb = b.data() + (size_t) row_b(i) * N;
|
||||
double me = 0, mr = 0;
|
||||
for (int o = 0; o < N; ++o) {
|
||||
if (!std::isfinite(ya[o])) throw std::runtime_error("a non-finite output");
|
||||
const double d = (double) ya[o] - yb[o];
|
||||
e2 += d * d;
|
||||
r2 += (double) yb[o] * yb[o];
|
||||
me = std::max(me, std::fabs(d));
|
||||
mr = std::max(mr, (double) std::fabs(yb[o]));
|
||||
}
|
||||
if (mr > 0) out.worst = std::max(out.worst, me / mr);
|
||||
else if (me > 0) out.worst = 1e30; // a zero row must stay exactly zero
|
||||
}
|
||||
out.rms = r2 > 0 ? std::sqrt(e2 / r2) : 0;
|
||||
return out;
|
||||
}
|
||||
|
||||
std::vector<float> download(const Dev& d, size_t n) {
|
||||
std::vector<float> h(n);
|
||||
ck(cudaMemcpy(h.data(), d.p, n * 4, cudaMemcpyDeviceToHost), "download");
|
||||
return h;
|
||||
}
|
||||
std::vector<int32_t> download_i(const Dev& d, size_t n) {
|
||||
std::vector<int32_t> h(n);
|
||||
ck(cudaMemcpy(h.data(), d.p, n * 4, cudaMemcpyDeviceToHost), "download");
|
||||
return h;
|
||||
}
|
||||
|
||||
// part 1: against the double-precision reference
|
||||
void reference_part(const Pair& p, cudaStream_t s) {
|
||||
constexpr int T = 256, E = 64, ZERO = 77;
|
||||
const Geo geo(p);
|
||||
std::mt19937 rng(136 + (unsigned) p.gu * 7 + (unsigned) p.d);
|
||||
std::vector<std::vector<uint8_t>> host((size_t) E);
|
||||
std::vector<std::unique_ptr<Dev>> dev;
|
||||
std::vector<const uint8_t*> blob((size_t) E);
|
||||
for (int e = 0; e < E; ++e) {
|
||||
host[(size_t) e] = make_blob(p, geo, rng);
|
||||
dev.push_back(std::make_unique<Dev>(geo.bytes + 4096 * (size_t) (e % 3) + 64)); // unrelated addresses
|
||||
uint8_t* q = dev.back()->as<uint8_t>() + 256 * (size_t) (e % 3) + 2 * (size_t) (e % 2); // 2-byte aligned
|
||||
ck(cudaMemcpy(q, host[(size_t) e].data(), geo.bytes, cudaMemcpyHostToDevice), "blob");
|
||||
blob[(size_t) e] = q;
|
||||
}
|
||||
const std::vector<float> x = make_x(T, rng, ZERO);
|
||||
const std::vector<int32_t> ids = make_ids(T, E, rng, 2, 4); // experts 0, 1 hot (several tiles), 60-63 unused
|
||||
const int64_t rows = (int64_t) T * K;
|
||||
const Routing r = sort_rows(ids, E);
|
||||
std::printf("%s - reference part: %d tokens, %d experts, rows per expert: max %d, experts without rows %d\n", p.name,
|
||||
T, E, *std::max_element(r.cnt.begin(), r.cnt.end()), (int) std::count(r.cnt.begin(), r.cnt.end(), 0));
|
||||
|
||||
Dev x_dev(x.size() * 4), src_dev((size_t) rows * 4);
|
||||
ck(cudaMemcpy(x_dev.p, x.data(), x.size() * 4, cudaMemcpyHostToDevice), "x");
|
||||
ck(cudaMemcpy(src_dev.p, r.src.data(), r.src.size() * 4, cudaMemcpyHostToDevice), "src");
|
||||
// (the uploads above are on the legacy stream and the work below on a non-blocking one: a cudaMemcpy from pageable
|
||||
// memory may return before its DMA lands - on Windows it can wait in the queue until the next synchronous call)
|
||||
ck(cudaDeviceSynchronize(), "uploads");
|
||||
mmq::Context ctx;
|
||||
MmqBufs mb(p, rows, E);
|
||||
run_mmq(p, geo, ctx, mb, r, x_dev.as<float>(), src_dev.as<int32_t>(), blob, rows, s);
|
||||
FusedBufs fb(T, rows, E);
|
||||
ck(cudaMemcpy(fb.ids.p, ids.data(), ids.size() * 4, cudaMemcpyHostToDevice), "ids");
|
||||
ck(cudaMemset(fb.dm.p, 0xff, (size_t) rows * N * 4), "sentinel"); // an unwritten row is NaN
|
||||
ck(cudaDeviceSynchronize(), "uploads");
|
||||
run_fused(p, geo, fb, x_dev.as<float>(), T, E, {0, 1, 20, E}, blob, s);
|
||||
ck(cudaStreamSynchronize(s), "sync");
|
||||
const std::vector<float> y_mmq = download(mb.dm, (size_t) rows * N), y_f = download(fb.dm, (size_t) rows * N);
|
||||
const std::vector<int32_t> slot = download_i(fb.slot, (size_t) rows), fsrc = download_i(fb.src, (size_t) rows);
|
||||
for (int64_t i = 0; i < rows; ++i)
|
||||
if (fsrc[(size_t) slot[(size_t) i]] != (int32_t) (i / K)) throw std::runtime_error("group: slot/src disagree");
|
||||
|
||||
// the reference, an expert at a time on all cores
|
||||
std::vector<float> ref((size_t) rows * N);
|
||||
std::vector<std::thread> th;
|
||||
const int nth = std::max(1, std::min(16, (int) std::thread::hardware_concurrency()));
|
||||
for (int w = 0; w < nth; ++w)
|
||||
th.emplace_back([&, w] {
|
||||
std::vector<float> gu, dn, h(FF);
|
||||
for (int e = w; e < E; e += nth) {
|
||||
if (r.cnt[(size_t) e] == 0) continue;
|
||||
dequant(p, geo, host[(size_t) e], gu, dn);
|
||||
for (int64_t i = 0; i < rows; ++i) {
|
||||
if (ids[(size_t) i] != e) continue;
|
||||
const float* xr = x.data() + (size_t) (i / K) * N;
|
||||
for (int f = 0; f < FF; ++f) {
|
||||
double gt = 0, up = 0;
|
||||
const float *wg = gu.data() + (size_t) f * N, *wu = gu.data() + (size_t) (FF + f) * N;
|
||||
for (int k = 0; k < N; ++k) { gt += (double) wg[k] * xr[k]; up += (double) wu[k] * xr[k]; }
|
||||
h[(size_t) f] = (float) (gt / (1.0 + std::exp(-gt)) * up);
|
||||
}
|
||||
for (int o = 0; o < N; ++o) {
|
||||
double a = 0;
|
||||
const float* wd = dn.data() + (size_t) o * FF;
|
||||
for (int f = 0; f < FF; ++f) a += (double) wd[f] * h[(size_t) f];
|
||||
ref[(size_t) i * N + o] = (float) a;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
for (auto& t : th) t.join();
|
||||
|
||||
auto pair = [](size_t i) { return (int64_t) i; };
|
||||
auto mrow = [&](size_t i) { return (int64_t) r.row_of[i]; };
|
||||
auto frow = [&](size_t i) { return (int64_t) slot[i]; };
|
||||
const Err em = compare(y_mmq, mrow, ref, pair, (size_t) rows);
|
||||
const Err ef = compare(y_f, frow, ref, pair, (size_t) rows);
|
||||
const Err efm = compare(y_f, frow, y_mmq, mrow, (size_t) rows);
|
||||
std::printf(" MMQ vs FP32 reference: rel RMS %.3e worst row max rel %.3e\n", em.rms, em.worst);
|
||||
std::printf(" fused vs FP32 reference: rel RMS %.3e worst row max rel %.3e\n", ef.rms, ef.worst);
|
||||
std::printf(" fused vs MMQ : rel RMS %.3e worst row max rel %.3e\n", efm.rms, efm.worst);
|
||||
double zmax = 0;
|
||||
for (int k = 0; k < K; ++k)
|
||||
for (int o = 0; o < N; ++o)
|
||||
zmax = std::max(zmax, (double) std::fabs(y_f[(size_t) slot[(size_t) (ZERO * K + k)] * N + o]));
|
||||
if (zmax != 0) throw std::runtime_error(std::string(p.name) + ": the all-zero token's outputs are not zero");
|
||||
if (ef.rms > 1.5 * em.rms || ef.worst > 2.0 * em.worst)
|
||||
throw std::runtime_error(std::string(p.name) + ": the fused path's error is not comparable to MMQ's");
|
||||
}
|
||||
|
||||
// part 2: one layer at a real chunk, timed
|
||||
void timing_part(const Pair& p, int T, cudaStream_t s) {
|
||||
constexpr int E = 512;
|
||||
const Geo geo(p);
|
||||
std::mt19937 rng(1360);
|
||||
const int64_t rows = (int64_t) T * K;
|
||||
// 512 distinct blobs - the products read each once per tile, as the engine does
|
||||
// at a stride of whole blocks of both formats (MMQ's expert stride counts blocks), + MMQ's read past the end
|
||||
size_t unit = ggml_type_size(p.gu);
|
||||
while (unit % ggml_type_size(p.d) != 0 || unit % 16 != 0) unit += ggml_type_size(p.gu);
|
||||
const size_t stride = (geo.bytes + unit - 1) / unit * unit;
|
||||
Dev blobs((size_t) E * stride + 4096);
|
||||
ck(cudaMemset(blobs.p, 0, (size_t) E * stride + 4096), "zero");
|
||||
{
|
||||
std::vector<uint8_t> b = make_blob(p, geo, rng);
|
||||
for (int e = 0; e < E; ++e) {
|
||||
b[4 + (size_t) (e % 16)] ^= (uint8_t) e; // not all equal (inside the first block, not its scale)
|
||||
ck(cudaMemcpy(blobs.as<uint8_t>() + (size_t) e * stride, b.data(), geo.bytes, cudaMemcpyHostToDevice),
|
||||
"blob");
|
||||
}
|
||||
}
|
||||
std::vector<const uint8_t*> blob((size_t) E);
|
||||
for (int e = 0; e < E; ++e) blob[(size_t) e] = blobs.as<uint8_t>() + (size_t) e * stride;
|
||||
const std::vector<float> x = make_x(T, rng, -1);
|
||||
const std::vector<int32_t> ids = make_ids(T, E, rng, 0, 0);
|
||||
const Routing r = sort_rows(ids, E);
|
||||
Dev x_dev(x.size() * 4), src_dev((size_t) rows * 4);
|
||||
ck(cudaMemcpy(x_dev.p, x.data(), x.size() * 4, cudaMemcpyHostToDevice), "x");
|
||||
ck(cudaMemcpy(src_dev.p, r.src.data(), r.src.size() * 4, cudaMemcpyHostToDevice), "src");
|
||||
mmq::Context ctx;
|
||||
MmqBufs mb(p, rows, E);
|
||||
FusedBufs fb(T, rows, E);
|
||||
ck(cudaMemcpy(fb.ids.p, ids.data(), ids.size() * 4, cudaMemcpyHostToDevice), "ids");
|
||||
ck(cudaDeviceSynchronize(), "uploads"); // (see reference_part)
|
||||
cudaEvent_t a, b;
|
||||
ck(cudaEventCreate(&a), "event");
|
||||
ck(cudaEventCreate(&b), "event");
|
||||
auto once = [&](auto&& f) {
|
||||
ck(cudaEventRecord(a, s), "record");
|
||||
f();
|
||||
ck(cudaEventRecord(b, s), "record");
|
||||
ck(cudaEventSynchronize(b), "sync");
|
||||
float t = 0;
|
||||
ck(cudaEventElapsedTime(&t, a, b), "elapsed");
|
||||
return t;
|
||||
};
|
||||
auto mmq_run = [&] { run_mmq(p, geo, ctx, mb, r, x_dev.as<float>(), src_dev.as<int32_t>(), blob, rows, s); };
|
||||
auto fused_run = [&] { run_fused(p, geo, fb, x_dev.as<float>(), T, E, {0, E}, blob, s); };
|
||||
const bool all_routed = std::count(r.cnt.begin(), r.cnt.end(), 0) == 0;
|
||||
auto direct_run = [&] {
|
||||
run_mmq(p, geo, ctx, mb, r, x_dev.as<float>(), src_dev.as<int32_t>(), blob, rows, s, true, stride);
|
||||
};
|
||||
if (std::getenv("S20_IDSDIAG")) { // which step overwrites the routing ids on the device
|
||||
auto check = [&](const char* after) {
|
||||
ck(cudaStreamSynchronize(s), "sync");
|
||||
const std::vector<int32_t> d = download_i(fb.ids, (size_t) rows);
|
||||
int64_t bad = 0;
|
||||
for (size_t i = 0; i < (size_t) rows; ++i) bad += d[i] != ids[i];
|
||||
std::printf(" ids on the device after %s: %lld of %lld differ (fb.ids %p, mb.dm %p..%p, fb.dm %p..%p)\n",
|
||||
after, (long long) bad, (long long) rows, fb.ids.p, mb.dm.p,
|
||||
(void*) (mb.dm.as<uint8_t>() + (size_t) rows * N * 4), fb.dm.p,
|
||||
(void*) (fb.dm.as<uint8_t>() + (size_t) rows * N * 4));
|
||||
};
|
||||
check("the upload");
|
||||
mmq_run();
|
||||
check("an MMQ run");
|
||||
fused_run();
|
||||
check("a fused run");
|
||||
if (all_routed) {
|
||||
direct_run();
|
||||
check("an MMQ run without gathers");
|
||||
}
|
||||
}
|
||||
// ~1 s of both first (the clocks ramp up), then alternating rounds; the medians
|
||||
const bool no_mmq = std::getenv("S20_NOMMQ") != nullptr; // (debug: the fused runs alone)
|
||||
for (float spent = 0; spent < 1000.0f;) spent += (no_mmq ? 0.0f : once(mmq_run)) + once(fused_run);
|
||||
std::vector<float> tm, tf, td;
|
||||
for (int rep = 0; rep < 8; ++rep) {
|
||||
tm.push_back(no_mmq ? 1.0f : once(mmq_run));
|
||||
tf.push_back(once(fused_run));
|
||||
if (all_routed && !no_mmq) td.push_back(once(direct_run));
|
||||
}
|
||||
if (td.empty()) td.push_back(0.0f);
|
||||
std::sort(tm.begin(), tm.end());
|
||||
std::sort(tf.begin(), tf.end());
|
||||
std::sort(td.begin(), td.end());
|
||||
const float t_mmq = tm[tm.size() / 2], t_f = tf[tf.size() / 2];
|
||||
if (all_routed)
|
||||
std::printf(" (MMQ's products alone, no gathers: %.3f ms - fused %.2fx)\n", td[td.size() / 2],
|
||||
td[td.size() / 2] / t_f);
|
||||
std::vector<float> tf2; // back to back
|
||||
for (int rep = 0; rep < 8; ++rep) tf2.push_back(once(fused_run));
|
||||
std::sort(tf2.begin(), tf2.end());
|
||||
std::printf(" (fused back to back: median %.3f ms, min %.3f; alternating min %.3f; MMQ min %.3f)\n", tf2[4],
|
||||
tf2[0], tf[0], tm[0]);
|
||||
const std::vector<float> y_mmq = download(mb.dm, (size_t) rows * N), y_f = download(fb.dm, (size_t) rows * N);
|
||||
const std::vector<int32_t> slot = download_i(fb.slot, (size_t) rows);
|
||||
const Err efm = compare(y_f, [&](size_t i) { return (int64_t) slot[i]; }, y_mmq,
|
||||
[&](size_t i) { return (int64_t) r.row_of[i]; }, (size_t) rows);
|
||||
std::printf("%s - timing part: %d tokens x top %d over %d experts, one layer: MMQ path %.3f ms, fused %.3f ms "
|
||||
"(%.2fx); fused vs MMQ rel RMS %.3e worst row %.3e\n",
|
||||
p.name, T, K, E, t_mmq, t_f, t_mmq / t_f, efm.rms, efm.worst);
|
||||
{ // the grouping: every pair in its expert's rows (the near-identical blobs here would hide a wrong expert)
|
||||
const std::vector<int32_t> fsrc = download_i(fb.src, (size_t) rows);
|
||||
int64_t bad = 0, bad_src = 0;
|
||||
for (size_t i = 0; i < (size_t) rows; ++i) {
|
||||
const int e = ids[i];
|
||||
bad += slot[i] < r.off[(size_t) e] || slot[i] >= r.off[(size_t) e + 1];
|
||||
bad_src += fsrc[(size_t) slot[i]] != (int32_t) (i / K);
|
||||
}
|
||||
if (bad || bad_src) {
|
||||
std::vector<int32_t> tbl((size_t) 4 * E + 2);
|
||||
ck(cudaMemcpy(tbl.data(), fb.scratch.p, tbl.size() * 4, cudaMemcpyDeviceToHost), "tables");
|
||||
int64_t bad_cnt = 0, bad_off = 0, bad_ids = 0;
|
||||
const std::vector<int32_t> d_ids = download_i(fb.ids, (size_t) rows);
|
||||
for (size_t i = 0; i < (size_t) rows; ++i) bad_ids += d_ids[i] != ids[i];
|
||||
std::printf(" the routing ids on the device: %lld of %lld differ (first %d %d %d)\n", (long long) bad_ids,
|
||||
(long long) rows, d_ids[0], d_ids[1], d_ids[2]);
|
||||
const int32_t smin = *std::min_element(slot.begin(), slot.end()), smax = *std::max_element(slot.begin(), slot.end());
|
||||
std::printf(" slots %d..%d; off[0..3] %d %d %d %d off[E] %d; fill[0..3] %d %d %d %d; ts[0..3] %d %d %d %d "
|
||||
"ts[E] %d\n", smin, smax, tbl[E], tbl[E + 1], tbl[E + 2], tbl[E + 3], tbl[2 * E], tbl[2 * E + 1],
|
||||
tbl[2 * E + 2], tbl[2 * E + 3], tbl[2 * E + 4], tbl[3 * E + 1], tbl[3 * E + 2], tbl[3 * E + 3],
|
||||
tbl[3 * E + 4], tbl[4 * E + 1]);
|
||||
// the group again on these ids, synchronously
|
||||
fused::group(fb.ids.as<int32_t>(), rows, K, E, fb.scratch.p, fb.slot.as<int32_t>(), fb.src.as<int32_t>(), s);
|
||||
ck(cudaStreamSynchronize(s), "sync");
|
||||
std::vector<int32_t> t2((size_t) 4 * E + 2);
|
||||
ck(cudaMemcpy(t2.data(), fb.scratch.p, t2.size() * 4, cudaMemcpyDeviceToHost), "tables");
|
||||
std::printf(" group() again, synchronized: cnt[0..3] %d %d %d %d off[E] %d\n", t2[0], t2[1], t2[2], t2[3],
|
||||
t2[(size_t) 2 * E]);
|
||||
for (int e = 0; e < E; ++e) {
|
||||
bad_cnt += tbl[(size_t) e] != r.cnt[(size_t) e];
|
||||
bad_off += tbl[(size_t) E + e] != r.off[(size_t) e];
|
||||
}
|
||||
std::printf(" GROUPING WRONG: %lld pairs outside their expert's rows, %lld src mismatches; tables: %lld counts "
|
||||
"and %lld offsets differ from the host's (cnt[0..3] %d %d %d %d vs %d %d %d %d)\n",
|
||||
(long long) bad, (long long) bad_src, (long long) bad_cnt, (long long) bad_off, tbl[0], tbl[1],
|
||||
tbl[2], tbl[3], r.cnt[0], r.cnt[1], r.cnt[2], r.cnt[3]);
|
||||
throw std::runtime_error(std::string(p.name) + ": the grouping is wrong");
|
||||
}
|
||||
}
|
||||
if (std::getenv("S20_ROWDIAG")) { // the pairs whose rows differ most: expert, its rows, the row's place
|
||||
const std::vector<int32_t> fsrc = download_i(fb.src, (size_t) rows);
|
||||
std::vector<std::pair<double, size_t>> worst;
|
||||
for (size_t i = 0; i < (size_t) rows; ++i) {
|
||||
const float* a = y_f.data() + (size_t) slot[i] * N;
|
||||
const float* m = y_mmq.data() + (size_t) r.row_of[i] * N;
|
||||
double me = 0, mr = 0;
|
||||
for (int o = 0; o < N; ++o) { me = std::max(me, (double) std::fabs(a[o] - m[o])); mr = std::max(mr, (double) std::fabs(m[o])); }
|
||||
worst.push_back({mr > 0 ? me / mr : me, i});
|
||||
}
|
||||
std::sort(worst.rbegin(), worst.rend());
|
||||
for (int k = 0; k < 12; ++k) {
|
||||
const size_t i = worst[(size_t) k].second;
|
||||
const int e = ids[i];
|
||||
// the expert's rows start at its first fused row: the smallest slot of its pairs
|
||||
int32_t first = INT32_MAX;
|
||||
for (size_t j = 0; j < (size_t) rows; ++j) if (ids[j] == e) first = std::min(first, slot[j]);
|
||||
std::printf(" pair %zu token %zu expert %d (%d rows) fused row %d (+%d in the expert) max rel %.3e\n", i,
|
||||
i / K, e, r.cnt[(size_t) e], slot[i], slot[i] - first, worst[(size_t) k].first);
|
||||
}
|
||||
(void) fsrc;
|
||||
}
|
||||
// both against the double-precision reference on a sample of 96 pairs
|
||||
{
|
||||
std::vector<size_t> pick;
|
||||
for (size_t i = 0; i < 96; ++i) pick.push_back((i * 2654435761u) % (size_t) rows);
|
||||
std::vector<float> ref(pick.size() * N), ym(pick.size() * N), yf(pick.size() * N);
|
||||
std::vector<uint8_t> hb(geo.bytes);
|
||||
std::vector<float> gu, dn, h(FF);
|
||||
for (size_t j = 0; j < pick.size(); ++j) {
|
||||
const size_t i = pick[j];
|
||||
const int e = ids[i];
|
||||
ck(cudaMemcpy(hb.data(), blob[(size_t) e], geo.bytes, cudaMemcpyDeviceToHost), "blob back");
|
||||
dequant(p, geo, hb, gu, dn);
|
||||
const float* xr = x.data() + (i / K) * N;
|
||||
for (int f = 0; f < FF; ++f) {
|
||||
double gt = 0, up = 0;
|
||||
const float *wg = gu.data() + (size_t) f * N, *wu = gu.data() + (size_t) (FF + f) * N;
|
||||
for (int k = 0; k < N; ++k) { gt += (double) wg[k] * xr[k]; up += (double) wu[k] * xr[k]; }
|
||||
h[(size_t) f] = (float) (gt / (1.0 + std::exp(-gt)) * up);
|
||||
}
|
||||
for (int o = 0; o < N; ++o) {
|
||||
double a = 0;
|
||||
const float* wd = dn.data() + (size_t) o * FF;
|
||||
for (int f = 0; f < FF; ++f) a += (double) wd[f] * h[(size_t) f];
|
||||
ref[j * N + o] = (float) a;
|
||||
}
|
||||
std::copy_n(y_mmq.begin() + (ptrdiff_t) ((size_t) r.row_of[i] * N), N, ym.begin() + (ptrdiff_t) (j * N));
|
||||
std::copy_n(y_f.begin() + (ptrdiff_t) ((size_t) slot[i] * N), N, yf.begin() + (ptrdiff_t) (j * N));
|
||||
}
|
||||
auto id = [](size_t j) { return (int64_t) j; };
|
||||
const Err em = compare(ym, id, ref, id, pick.size()), ef = compare(yf, id, ref, id, pick.size());
|
||||
std::printf(" sample of %zu pairs vs FP32 reference: MMQ rel RMS %.3e worst %.3e, fused rel RMS %.3e worst "
|
||||
"%.3e\n", pick.size(), em.rms, em.worst, ef.rms, ef.worst);
|
||||
if (ef.rms > 1.5 * em.rms || ef.worst > 2.0 * em.worst)
|
||||
throw std::runtime_error(std::string(p.name) + ": the fused path's error is not comparable to MMQ's (timing)");
|
||||
}
|
||||
cudaEventDestroy(a);
|
||||
cudaEventDestroy(b);
|
||||
if (efm.rms > 0.05) throw std::runtime_error(std::string(p.name) + ": the paths disagree at the real shape");
|
||||
}
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
try {
|
||||
#ifdef _WIN32
|
||||
_putenv_s("STRATA_PF_FUSED", "1");
|
||||
#else
|
||||
setenv("STRATA_PF_FUSED", "1", 1);
|
||||
#endif
|
||||
int n = 0;
|
||||
if (cudaGetDeviceCount(&n) != cudaSuccess || n == 0) { std::printf("no CUDA device: skipped\n"); return 77; }
|
||||
if (!fused::available()) { std::printf("the fused kernels need sm_80 or newer: skipped\n"); return 77; }
|
||||
if (!mmq::built()) { std::printf("no MMQ in this build\n"); return 1; }
|
||||
const std::vector<Pair> pairs = {
|
||||
{GGML_TYPE_IQ2_S, GGML_TYPE_Q2_0, "IQ2_S / Q2_0"}, {GGML_TYPE_IQ2_XXS, GGML_TYPE_Q2_0, "IQ2_XXS / Q2_0"},
|
||||
{GGML_TYPE_IQ2_XS, GGML_TYPE_IQ4_NL, "IQ2_XS / IQ4_NL"}, {GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ4_NL, "IQ3_XXS / IQ4_NL"},
|
||||
{GGML_TYPE_IQ3_S, GGML_TYPE_IQ4_NL, "IQ3_S / IQ4_NL"}, {GGML_TYPE_IQ4_XS, GGML_TYPE_Q2_0, "IQ4_XS / Q2_0"}};
|
||||
const char* only = nullptr; // --only=NAME: the pairs whose name contains NAME
|
||||
for (int i = 1; i < argc; ++i)
|
||||
if (std::strncmp(argv[i], "--only=", 7) == 0) only = argv[i] + 7;
|
||||
bool timing = true, ref = true; // --no-timing, --no-ref: skip a part
|
||||
std::vector<int> chunks = {2048, 3584, 8192}; // --chunks=A,B: the timing part's chunk sizes
|
||||
for (int i = 1; i < argc; ++i) {
|
||||
if (std::strcmp(argv[i], "--no-timing") == 0) timing = false;
|
||||
if (std::strcmp(argv[i], "--no-ref") == 0) ref = false;
|
||||
if (std::strncmp(argv[i], "--chunks=", 9) == 0) {
|
||||
chunks.clear();
|
||||
for (const char* c = argv[i] + 9; *c;) {
|
||||
chunks.push_back(std::atoi(c));
|
||||
while (*c && *c != ',') ++c;
|
||||
if (*c == ',') ++c;
|
||||
}
|
||||
}
|
||||
}
|
||||
cudaStream_t s = nullptr;
|
||||
ck(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking), "stream");
|
||||
for (const Pair& p : pairs) {
|
||||
if (only && std::string(p.name).find(only) == std::string::npos) continue;
|
||||
if (!fused::native_supported((int) p.gu, (int) p.d))
|
||||
throw std::runtime_error(std::string(p.name) + ": not covered by the native kernels");
|
||||
if (ref) reference_part(p, s);
|
||||
}
|
||||
if (timing) {
|
||||
// the IQ2_XS pack's layers (34 of 48 IQ2_S / Q2_0, 11 IQ2_XXS / Q2_0), the IQ3_S pack's most common, IQ3_XXS's
|
||||
for (const Pair& p : {pairs[0], pairs[1], pairs[4], pairs[3]}) {
|
||||
if (only && std::string(p.name).find(only) == std::string::npos) continue;
|
||||
for (int T : chunks) timing_part(p, T, s);
|
||||
}
|
||||
}
|
||||
ck(cudaStreamDestroy(s), "destroy");
|
||||
std::printf("prefill fused MoE (native formats) parity passed\n");
|
||||
return 0;
|
||||
} catch (const std::exception& e) {
|
||||
std::fprintf(stderr, "prefill fused MoE (native formats) parity failed: %s\n", e.what());
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,420 @@
|
||||
// prefill_fused_moe_test - #136: the fused int8 expert kernels of the prompt path (moe_fused.hpp, STRATA_PF_FUSED=1)
|
||||
// on random Strata Q2_0 expert blobs and random routing, against
|
||||
// - a double-precision reference: the dequantized weights times the FP32 activations (gate/up, SwiGLU, down), and
|
||||
// - the MMQ path in prefill.cpp's sequence (gather_strata_q2 into 16-expert groups, q8_1 of the slots' activations,
|
||||
// gate/up, SwiGLU, q8_1 of H, down).
|
||||
// Both round the activations and H to int8 per 32 values, differently, so this is not a bitwise test: the fused
|
||||
// path's error against the reference has to be comparable to MMQ's own (at most 1.5x its RMS and 2x its worst row).
|
||||
// Part 1 (reference): 256 tokens over 64 experts, skewed routing - experts with 0 rows, with one, with several 64-row
|
||||
// tiles - an all-zero token, per-expert blobs at unrelated addresses, three launches of different expert ranges.
|
||||
// Part 2 (timing): one layer at a real chunk (2048 tokens, 512 experts, top 10) on both paths, with their agreement.
|
||||
// Exit 77 without a CUDA device of sm_80 or newer.
|
||||
#include "strata/prefill/moe_fused.hpp"
|
||||
#include "strata/prefill/moe_mmq.hpp"
|
||||
|
||||
#include "ggml.h"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <random>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
namespace mmq = strata::prefill::mmq;
|
||||
namespace fused = strata::prefill::fused;
|
||||
|
||||
constexpr int N = 2560, FF = 640, K = 10, GROUP = 16, Q2 = 42; // GGML_TYPE_Q2_0
|
||||
constexpr size_t BLOB = 1382400, O_D_CODES = (size_t) 1280 * 640, O_GU_SC = O_D_CODES + (size_t) 2560 * 160,
|
||||
O_D_SC = O_GU_SC + (size_t) 1280 * 40 * 2;
|
||||
|
||||
void ck(cudaError_t e, const char* what) {
|
||||
if (e != cudaSuccess) throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(e));
|
||||
}
|
||||
|
||||
struct Dev {
|
||||
void* p = nullptr;
|
||||
explicit Dev(size_t n) { ck(cudaMalloc(&p, n), "cudaMalloc"); }
|
||||
~Dev() { cudaFree(p); }
|
||||
Dev(const Dev&) = delete;
|
||||
Dev& operator=(const Dev&) = delete;
|
||||
template <typename T> T* as() const { return (T*) p; }
|
||||
};
|
||||
|
||||
// a Strata Q2_0 blob: random codes, scales in [0.004, 0.03] (a 2-bit expert's range)
|
||||
std::vector<uint8_t> make_blob(std::mt19937& rng) {
|
||||
std::vector<uint8_t> b(BLOB);
|
||||
std::uniform_int_distribution<int> byte(0, 255);
|
||||
std::uniform_real_distribution<float> sc(0.004f, 0.03f);
|
||||
for (size_t i = 0; i < O_GU_SC; ++i) b[i] = (uint8_t) byte(rng);
|
||||
for (size_t i = O_GU_SC; i < BLOB; i += 2) {
|
||||
const ggml_fp16_t h = ggml_fp32_to_fp16(sc(rng));
|
||||
std::memcpy(&b[i], &h, 2);
|
||||
}
|
||||
return b;
|
||||
}
|
||||
|
||||
// the blob's weights as floats: gate/up [1280][2560] (interleaved rows), down [2560][640]
|
||||
void dequant(const std::vector<uint8_t>& b, std::vector<float>& gu, std::vector<float>& dn) {
|
||||
auto row = [&](size_t codes, size_t ld, size_t scales, int nb, int r, float* out) {
|
||||
for (int k = 0; k < nb * 64; ++k) {
|
||||
ggml_fp16_t h;
|
||||
std::memcpy(&h, &b[scales + ((size_t) r * nb + k / 64) * 2], 2);
|
||||
const int q = (b[codes + (size_t) r * ld + k / 4] >> (2 * (k % 4))) & 3;
|
||||
out[k] = ggml_fp16_to_fp32(h) * (float) (q - 1);
|
||||
}
|
||||
};
|
||||
gu.resize((size_t) 1280 * N);
|
||||
dn.resize((size_t) N * FF);
|
||||
for (int r = 0; r < 1280; ++r) row(0, 640, O_GU_SC, 40, r, gu.data() + (size_t) r * N);
|
||||
for (int r = 0; r < N; ++r) row(O_D_CODES, 160, O_D_SC, 10, r, dn.data() + (size_t) r * FF);
|
||||
}
|
||||
|
||||
// activations: N(0, 1) with a few large values (the hidden state has outliers); token `zero` all zero
|
||||
std::vector<float> make_x(int T, std::mt19937& rng, int zero) {
|
||||
std::vector<float> x((size_t) T * N);
|
||||
std::normal_distribution<float> nd(0.0f, 1.0f);
|
||||
std::uniform_int_distribution<int> pick(0, 199);
|
||||
for (float& v : x) v = nd(rng) * (pick(rng) == 0 ? 12.0f : 1.0f);
|
||||
if (zero >= 0) std::fill(x.begin() + (size_t) zero * N, x.begin() + (size_t) (zero + 1) * N, 0.0f);
|
||||
return x;
|
||||
}
|
||||
|
||||
// K distinct experts per token; `hot` > 0 skews the first choices towards the lowest ids
|
||||
std::vector<int32_t> make_ids(int T, int E, std::mt19937& rng, int hot, int unused) {
|
||||
std::vector<int32_t> ids((size_t) T * K);
|
||||
std::uniform_int_distribution<int> any(0, E - 1 - unused);
|
||||
for (int t = 0; t < T; ++t)
|
||||
for (int k = 0; k < K; ++k) {
|
||||
int e;
|
||||
bool dup;
|
||||
do {
|
||||
e = (hot > 0 && k < 3 && (int) (rng() % 4) != 0) ? (int) (rng() % (unsigned) hot) : any(rng);
|
||||
dup = false;
|
||||
for (int j = 0; j < k; ++j) dup |= ids[(size_t) t * K + j] == e;
|
||||
} while (dup);
|
||||
ids[(size_t) t * K + k] = e;
|
||||
}
|
||||
return ids;
|
||||
}
|
||||
|
||||
struct Routing {
|
||||
std::vector<int32_t> cnt, off, src; // per expert; sorted row -> token (rows in pair order within an expert)
|
||||
std::vector<int32_t> row_of; // pair -> sorted row
|
||||
};
|
||||
Routing sort_rows(const std::vector<int32_t>& ids, int E) {
|
||||
Routing r;
|
||||
r.cnt.assign((size_t) E, 0);
|
||||
for (int32_t e : ids) ++r.cnt[(size_t) e];
|
||||
r.off.assign((size_t) E + 1, 0);
|
||||
for (int e = 0; e < E; ++e) r.off[(size_t) e + 1] = r.off[(size_t) e] + r.cnt[(size_t) e];
|
||||
std::vector<int32_t> fill(r.off.begin(), r.off.end() - 1);
|
||||
r.src.resize(ids.size());
|
||||
r.row_of.resize(ids.size());
|
||||
for (size_t i = 0; i < ids.size(); ++i) {
|
||||
const int32_t p = fill[(size_t) ids[i]]++;
|
||||
r.row_of[i] = p;
|
||||
r.src[(size_t) p] = (int32_t) (i / K);
|
||||
}
|
||||
return r;
|
||||
}
|
||||
|
||||
// MMQ, as prefill.cpp runs it: per group of 16 routed experts a gather, gate/up, SwiGLU, q8_1 of H, down. Returns
|
||||
// Dm [rows][N] in sorted-row order.
|
||||
struct MmqBufs {
|
||||
Dev xq, gu, h, hq, dm, ident, bounds, grp_gu, grp_d;
|
||||
MmqBufs(int64_t rows, int E)
|
||||
: xq(mmq::q8_bytes(rows, N)), gu((size_t) rows * 1280 * 4), h((size_t) rows * FF * 4),
|
||||
hq(mmq::q8_bytes(rows, FF)),
|
||||
dm((size_t) rows * N * 4), ident((size_t) rows * 4), bounds((size_t) (2 * (E + E / GROUP + 2)) * 4),
|
||||
grp_gu(GROUP * mmq::matrix_bytes(Q2, 1280, N) + 4096), grp_d(GROUP * mmq::matrix_bytes(Q2, N, FF) + 4096) {}
|
||||
};
|
||||
void run_mmq(mmq::Context& ctx, MmqBufs& b, const Routing& r, const float* x_dev, const int32_t* src_dev,
|
||||
const std::vector<const uint8_t*>& blob, int64_t rows, cudaStream_t s) {
|
||||
const size_t gub = mmq::matrix_bytes(Q2, 1280, N), db = mmq::matrix_bytes(Q2, N, FF);
|
||||
const int E = (int) r.cnt.size();
|
||||
std::vector<int32_t> order;
|
||||
for (int e = 0; e < E; ++e) if (r.cnt[(size_t) e] > 0) order.push_back(e);
|
||||
const size_t n = order.size(), ng = (n + GROUP - 1) / GROUP;
|
||||
std::vector<int32_t> bh(n + 1 + ng * (GROUP + 1));
|
||||
for (size_t j = 0; j < n; ++j) bh[j] = r.off[(size_t) order[j]];
|
||||
bh[n] = (int32_t) rows;
|
||||
for (size_t g = 0; g < ng; ++g)
|
||||
for (size_t i = 0; i <= GROUP; ++i)
|
||||
bh[n + 1 + g * (GROUP + 1) + i] = bh[std::min(n, g * GROUP + i)] - bh[g * GROUP];
|
||||
ck(cudaMemcpyAsync(b.bounds.p, bh.data(), bh.size() * 4, cudaMemcpyHostToDevice, s), "bounds");
|
||||
mmq::iota(b.ident.as<int32_t>(), rows, s);
|
||||
mmq::quantize(x_dev, src_dev, b.xq.p, Q2, N, N, rows, s);
|
||||
for (size_t j = 0; j < n; ++j) {
|
||||
const size_t q = j % GROUP;
|
||||
mmq::gather_strata_q2(blob[(size_t) order[j]], b.grp_gu.as<uint8_t>() + q * gub, b.grp_d.as<uint8_t>() + q * db,
|
||||
s);
|
||||
if (q + 1 < GROUP && j + 1 < n) continue;
|
||||
const size_t j0 = j - q, g = j0 / GROUP;
|
||||
const int ngx = (int) (q + 1);
|
||||
const int64_t r0 = bh[j0], nr = bh[j + 1] - r0;
|
||||
int64_t maxr = 0;
|
||||
for (size_t i = j0; i <= j; ++i) maxr = std::max<int64_t>(maxr, r.cnt[(size_t) order[i]]);
|
||||
ck(cudaMemsetAsync(b.grp_gu.as<uint8_t>() + (size_t) ngx * gub, 0, 4096, s), "tail");
|
||||
ck(cudaMemsetAsync(b.grp_d.as<uint8_t>() + (size_t) ngx * db, 0, 4096, s), "tail");
|
||||
mmq::Product gu;
|
||||
gu.w = b.grp_gu.p; gu.type = Q2; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = gub;
|
||||
gu.n = ngx; gu.xq = b.xq.p; gu.bounds = b.bounds.as<int32_t>() + j0; gu.ids = b.ident.as<int32_t>();
|
||||
gu.total_rows = rows; gu.max_rows = maxr; gu.dst = b.gu.as<float>(); gu.ld_dst = 1280;
|
||||
ctx.run(gu, s);
|
||||
mmq::swiglu(b.gu.as<float>() + r0 * 1280, b.h.as<float>() + r0 * FF, nr, FF, true, s);
|
||||
mmq::quantize(b.h.as<float>() + r0 * FF, nullptr, b.hq.p, Q2, FF, FF, nr, s);
|
||||
mmq::Product dn;
|
||||
dn.w = b.grp_d.p; dn.type = Q2; dn.w_rows = N; dn.w_cols = FF; dn.expert_bytes = db;
|
||||
dn.n = ngx; dn.xq = b.hq.p; dn.bounds = b.bounds.as<int32_t>() + n + 1 + g * (GROUP + 1);
|
||||
dn.ids = b.ident.as<int32_t>(); dn.total_rows = nr; dn.max_rows = maxr; dn.dst = b.dm.as<float>() + r0 * N;
|
||||
dn.ld_dst = N;
|
||||
ctx.run(dn, s);
|
||||
}
|
||||
}
|
||||
|
||||
// the fused path: quantize, group on the device, then launches over expert ranges `cuts` (e.g. {0, 20, 64})
|
||||
struct FusedBufs {
|
||||
Dev xa, ha, dm, scratch, slot, src, ids;
|
||||
FusedBufs(int T, int64_t rows, int E)
|
||||
: xa(fused::act_bytes(T, N)), ha(fused::act_bytes(rows, FF)), dm((size_t) rows * N * 4),
|
||||
scratch(fused::group_bytes(rows, E)), slot((size_t) rows * 4), src((size_t) rows * 4),
|
||||
ids((size_t) rows * 4) {}
|
||||
};
|
||||
void run_fused(FusedBufs& b, const float* x_dev, int T, int E, const std::vector<int>& cuts,
|
||||
const std::vector<const uint8_t*>& blob, cudaStream_t s) {
|
||||
const int64_t rows = (int64_t) T * K;
|
||||
fused::quantize_act(x_dev, T, N, b.xa.p, s);
|
||||
fused::group(b.ids.as<int32_t>(), rows, K, E, b.scratch.p, b.slot.as<int32_t>(), b.src.as<int32_t>(), s);
|
||||
for (size_t c = 0; c + 1 < cuts.size(); ++c) {
|
||||
for (int e0 = cuts[c]; e0 < cuts[c + 1]; e0 += fused::kMaxBatch) {
|
||||
fused::Batch bt;
|
||||
bt.e0 = e0;
|
||||
bt.e1 = std::min(cuts[c + 1], e0 + fused::kMaxBatch);
|
||||
for (int e = bt.e0; e < bt.e1; ++e) bt.blob[e - bt.e0] = blob[(size_t) e];
|
||||
fused::experts(bt, E, rows, b.scratch.p, b.xa.p, b.src.as<int32_t>(), b.ha.p, b.dm.as<float>(), s);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct Err {
|
||||
double rms = 0, worst = 0; // RMS error / RMS of the reference; the worst row's max |error| / its max |ref|
|
||||
};
|
||||
// y_a[row_a(i)] against y_b[row_b(i)] for every pair i (rows of N values)
|
||||
template <typename RA, typename RB>
|
||||
Err compare(const std::vector<float>& a, RA row_a, const std::vector<float>& b, RB row_b, size_t pairs) {
|
||||
double e2 = 0, r2 = 0;
|
||||
Err out;
|
||||
for (size_t i = 0; i < pairs; ++i) {
|
||||
const float* ya = a.data() + (size_t) row_a(i) * N;
|
||||
const float* yb = b.data() + (size_t) row_b(i) * N;
|
||||
double me = 0, mr = 0;
|
||||
for (int o = 0; o < N; ++o) {
|
||||
if (!std::isfinite(ya[o])) throw std::runtime_error("a non-finite output");
|
||||
const double d = (double) ya[o] - yb[o];
|
||||
e2 += d * d;
|
||||
r2 += (double) yb[o] * yb[o];
|
||||
me = std::max(me, std::fabs(d));
|
||||
mr = std::max(mr, (double) std::fabs(yb[o]));
|
||||
}
|
||||
if (mr > 0) out.worst = std::max(out.worst, me / mr);
|
||||
else if (me > 0) out.worst = 1e30; // a zero row must stay exactly zero
|
||||
}
|
||||
out.rms = r2 > 0 ? std::sqrt(e2 / r2) : 0;
|
||||
return out;
|
||||
}
|
||||
|
||||
std::vector<float> download(const Dev& d, size_t n) {
|
||||
std::vector<float> h(n);
|
||||
ck(cudaMemcpy(h.data(), d.p, n * 4, cudaMemcpyDeviceToHost), "download");
|
||||
return h;
|
||||
}
|
||||
std::vector<int32_t> download_i(const Dev& d, size_t n) {
|
||||
std::vector<int32_t> h(n);
|
||||
ck(cudaMemcpy(h.data(), d.p, n * 4, cudaMemcpyDeviceToHost), "download");
|
||||
return h;
|
||||
}
|
||||
|
||||
// part 1: against the double-precision reference
|
||||
void reference_part(cudaStream_t s) {
|
||||
constexpr int T = 256, E = 64, ZERO = 77;
|
||||
std::mt19937 rng(136);
|
||||
std::vector<std::vector<uint8_t>> host((size_t) E);
|
||||
std::vector<std::unique_ptr<Dev>> dev;
|
||||
std::vector<const uint8_t*> blob((size_t) E);
|
||||
for (int e = 0; e < E; ++e) {
|
||||
host[(size_t) e] = make_blob(rng);
|
||||
dev.push_back(std::make_unique<Dev>(BLOB + 4096 * (size_t) (e % 3))); // unrelated addresses
|
||||
uint8_t* p = dev.back()->as<uint8_t>() + 256 * (size_t) (e % 3); // 16-byte aligned, not 512
|
||||
ck(cudaMemcpy(p, host[(size_t) e].data(), BLOB, cudaMemcpyHostToDevice), "blob");
|
||||
blob[(size_t) e] = p;
|
||||
}
|
||||
const std::vector<float> x = make_x(T, rng, ZERO);
|
||||
const std::vector<int32_t> ids = make_ids(T, E, rng, 2, 4); // experts 0, 1 hot (several tiles), 60-63 unused
|
||||
const int64_t rows = (int64_t) T * K;
|
||||
const Routing r = sort_rows(ids, E);
|
||||
std::printf("reference part: %d tokens, %d experts, rows per expert: max %d, experts without rows %d\n", T, E,
|
||||
*std::max_element(r.cnt.begin(), r.cnt.end()), (int) std::count(r.cnt.begin(), r.cnt.end(), 0));
|
||||
|
||||
Dev x_dev(x.size() * 4), src_dev((size_t) rows * 4);
|
||||
ck(cudaMemcpy(x_dev.p, x.data(), x.size() * 4, cudaMemcpyHostToDevice), "x");
|
||||
ck(cudaMemcpy(src_dev.p, r.src.data(), r.src.size() * 4, cudaMemcpyHostToDevice), "src");
|
||||
mmq::Context ctx;
|
||||
MmqBufs mb(rows, E);
|
||||
// the uploads (legacy stream) must land before the non-blocking stream reads them (S20 found the kernels could run
|
||||
// on all-zero routing otherwise - every pair on one expert: fast and wrong)
|
||||
ck(cudaDeviceSynchronize(), "uploads");
|
||||
run_mmq(ctx, mb, r, x_dev.as<float>(), src_dev.as<int32_t>(), blob, rows, s);
|
||||
FusedBufs fb(T, rows, E);
|
||||
ck(cudaMemcpy(fb.ids.p, ids.data(), ids.size() * 4, cudaMemcpyHostToDevice), "ids");
|
||||
ck(cudaMemset(fb.dm.p, 0xff, (size_t) rows * N * 4), "sentinel"); // an unwritten row is NaN
|
||||
ck(cudaDeviceSynchronize(), "uploads");
|
||||
run_fused(fb, x_dev.as<float>(), T, E, {0, 1, 20, E}, blob, s);
|
||||
ck(cudaStreamSynchronize(s), "sync");
|
||||
const std::vector<float> y_mmq = download(mb.dm, (size_t) rows * N), y_f = download(fb.dm, (size_t) rows * N);
|
||||
const std::vector<int32_t> slot = download_i(fb.slot, (size_t) rows), fsrc = download_i(fb.src, (size_t) rows);
|
||||
for (int64_t i = 0; i < rows; ++i)
|
||||
if (fsrc[(size_t) slot[(size_t) i]] != (int32_t) (i / K)) throw std::runtime_error("group: slot/src disagree");
|
||||
|
||||
// the reference, an expert at a time on all cores
|
||||
std::vector<float> ref((size_t) rows * N);
|
||||
std::vector<std::thread> th;
|
||||
const int nth = std::max(1, std::min(16, (int) std::thread::hardware_concurrency()));
|
||||
for (int w = 0; w < nth; ++w)
|
||||
th.emplace_back([&, w] {
|
||||
std::vector<float> gu, dn, h(FF);
|
||||
for (int e = w; e < E; e += nth) {
|
||||
if (r.cnt[(size_t) e] == 0) continue;
|
||||
dequant(host[(size_t) e], gu, dn);
|
||||
for (int64_t i = 0; i < rows; ++i) {
|
||||
if (ids[(size_t) i] != e) continue;
|
||||
const float* xr = x.data() + (size_t) (i / K) * N;
|
||||
for (int f = 0; f < FF; ++f) {
|
||||
double gt = 0, up = 0;
|
||||
const float *wg = gu.data() + (size_t) (2 * f) * N, *wu = wg + N;
|
||||
for (int k = 0; k < N; ++k) { gt += (double) wg[k] * xr[k]; up += (double) wu[k] * xr[k]; }
|
||||
h[(size_t) f] = (float) (gt / (1.0 + std::exp(-gt)) * up);
|
||||
}
|
||||
for (int o = 0; o < N; ++o) {
|
||||
double a = 0;
|
||||
const float* wd = dn.data() + (size_t) o * FF;
|
||||
for (int f = 0; f < FF; ++f) a += (double) wd[f] * h[(size_t) f];
|
||||
ref[(size_t) i * N + o] = (float) a;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
for (auto& t : th) t.join();
|
||||
|
||||
auto pair = [](size_t i) { return (int64_t) i; };
|
||||
auto mrow = [&](size_t i) { return (int64_t) r.row_of[i]; };
|
||||
auto frow = [&](size_t i) { return (int64_t) slot[i]; };
|
||||
const Err em = compare(y_mmq, mrow, ref, pair, (size_t) rows);
|
||||
const Err ef = compare(y_f, frow, ref, pair, (size_t) rows);
|
||||
const Err efm = compare(y_f, frow, y_mmq, mrow, (size_t) rows);
|
||||
std::printf(" MMQ vs FP32 reference: rel RMS %.3e worst row max rel %.3e\n", em.rms, em.worst);
|
||||
std::printf(" fused vs FP32 reference: rel RMS %.3e worst row max rel %.3e\n", ef.rms, ef.worst);
|
||||
std::printf(" fused vs MMQ : rel RMS %.3e worst row max rel %.3e\n", efm.rms, efm.worst);
|
||||
double zmax = 0;
|
||||
for (int k = 0; k < K; ++k)
|
||||
for (int o = 0; o < N; ++o)
|
||||
zmax = std::max(zmax, (double) std::fabs(y_f[(size_t) slot[(size_t) (ZERO * K + k)] * N + o]));
|
||||
if (zmax != 0) throw std::runtime_error("the all-zero token's outputs are not zero");
|
||||
if (ef.rms > 1.5 * em.rms || ef.worst > 2.0 * em.worst)
|
||||
throw std::runtime_error("the fused path's error is not comparable to MMQ's");
|
||||
}
|
||||
|
||||
// part 2: one layer at a real chunk, timed
|
||||
void timing_part(cudaStream_t s) {
|
||||
constexpr int T = 2048, E = 512;
|
||||
std::mt19937 rng(1360);
|
||||
const int64_t rows = (int64_t) T * K;
|
||||
// 512 distinct blobs (708 MB) - the products read each once per tile, as the engine does
|
||||
Dev blobs((size_t) E * BLOB);
|
||||
{
|
||||
std::vector<uint8_t> b = make_blob(rng);
|
||||
for (int e = 0; e < E; ++e) {
|
||||
b[(size_t) e * 7919 % O_GU_SC] ^= (uint8_t) e; // not all equal
|
||||
ck(cudaMemcpy(blobs.as<uint8_t>() + (size_t) e * BLOB, b.data(), BLOB, cudaMemcpyHostToDevice), "blob");
|
||||
}
|
||||
}
|
||||
std::vector<const uint8_t*> blob((size_t) E);
|
||||
for (int e = 0; e < E; ++e) blob[(size_t) e] = blobs.as<uint8_t>() + (size_t) e * BLOB;
|
||||
const std::vector<float> x = make_x(T, rng, -1);
|
||||
const std::vector<int32_t> ids = make_ids(T, E, rng, 0, 0);
|
||||
const Routing r = sort_rows(ids, E);
|
||||
Dev x_dev(x.size() * 4), src_dev((size_t) rows * 4);
|
||||
ck(cudaMemcpy(x_dev.p, x.data(), x.size() * 4, cudaMemcpyHostToDevice), "x");
|
||||
ck(cudaMemcpy(src_dev.p, r.src.data(), r.src.size() * 4, cudaMemcpyHostToDevice), "src");
|
||||
mmq::Context ctx;
|
||||
MmqBufs mb(rows, E);
|
||||
FusedBufs fb(T, rows, E);
|
||||
ck(cudaMemcpy(fb.ids.p, ids.data(), ids.size() * 4, cudaMemcpyHostToDevice), "ids");
|
||||
ck(cudaDeviceSynchronize(), "uploads"); // the legacy-stream uploads before the timed non-blocking stream
|
||||
cudaEvent_t a, b;
|
||||
ck(cudaEventCreate(&a), "event");
|
||||
ck(cudaEventCreate(&b), "event");
|
||||
auto time = [&](auto&& f) {
|
||||
std::vector<float> ms;
|
||||
for (int rep = 0; rep < 6; ++rep) {
|
||||
ck(cudaEventRecord(a, s), "record");
|
||||
f();
|
||||
ck(cudaEventRecord(b, s), "record");
|
||||
ck(cudaEventSynchronize(b), "sync");
|
||||
float t = 0;
|
||||
ck(cudaEventElapsedTime(&t, a, b), "elapsed");
|
||||
if (rep > 0) ms.push_back(t); // the first is the warm-up
|
||||
}
|
||||
std::sort(ms.begin(), ms.end());
|
||||
return ms[ms.size() / 2];
|
||||
};
|
||||
// the engine's batches: up to 128 experts per launch (the streamed ring decides; 4 launches here)
|
||||
const float t_mmq = time([&] { run_mmq(ctx, mb, r, x_dev.as<float>(), src_dev.as<int32_t>(), blob, rows, s); });
|
||||
const float t_f = time([&] { run_fused(fb, x_dev.as<float>(), T, E, {0, E}, blob, s); });
|
||||
const float t_f1 = time([&] { // one launch per 16 experts, MMQ's group size
|
||||
std::vector<int> cuts;
|
||||
for (int e = 0; e <= E; e += 16) cuts.push_back(e);
|
||||
run_fused(fb, x_dev.as<float>(), T, E, cuts, blob, s);
|
||||
});
|
||||
const std::vector<float> y_mmq = download(mb.dm, (size_t) rows * N), y_f = download(fb.dm, (size_t) rows * N);
|
||||
const std::vector<int32_t> slot = download_i(fb.slot, (size_t) rows);
|
||||
const Err efm = compare(y_f, [&](size_t i) { return (int64_t) slot[i]; }, y_mmq,
|
||||
[&](size_t i) { return (int64_t) r.row_of[i]; }, (size_t) rows);
|
||||
std::printf("timing part: %d tokens x top %d over %d experts (%.1f rows per expert), one layer:\n", T, K, E,
|
||||
(double) rows / E);
|
||||
std::printf(" MMQ path (gathers, q8_1, products, SwiGLU, 32 groups): %.3f ms\n", t_mmq);
|
||||
std::printf(" fused path (quantize, group, 4 launches of 128): %.3f ms (%.2fx)\n", t_f, t_mmq / t_f);
|
||||
std::printf(" fused path, launches of 16 experts: %.3f ms\n", t_f1);
|
||||
std::printf(" fused vs MMQ: rel RMS %.3e worst row max rel %.3e\n", efm.rms, efm.worst);
|
||||
cudaEventDestroy(a);
|
||||
cudaEventDestroy(b);
|
||||
if (efm.rms > 0.05) throw std::runtime_error("the paths disagree at the real shape");
|
||||
}
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
try {
|
||||
int n = 0;
|
||||
if (cudaGetDeviceCount(&n) != cudaSuccess || n == 0) { std::printf("no CUDA device: skipped\n"); return 77; }
|
||||
if (!fused::available()) { std::printf("the fused kernels need sm_80 or newer: skipped\n"); return 77; }
|
||||
if (!mmq::built()) { std::printf("no MMQ in this build\n"); return 1; }
|
||||
cudaStream_t s = nullptr;
|
||||
ck(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking), "stream");
|
||||
reference_part(s);
|
||||
if (!(argc > 1 && std::strcmp(argv[1], "--no-timing") == 0)) timing_part(s);
|
||||
ck(cudaStreamDestroy(s), "destroy");
|
||||
std::printf("prefill fused MoE parity passed\n");
|
||||
return 0;
|
||||
} catch (const std::exception& e) {
|
||||
std::fprintf(stderr, "prefill fused MoE parity failed: %s\n", e.what());
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
+24
-4
@@ -119,10 +119,30 @@ def run(cfg: dict, say=print, start_engine=None) -> dict:
|
||||
tok = ST.Tokenizer(toks, (tpath / "merges.txt").read_text(encoding="utf-8").split("\n"),
|
||||
json.loads((tpath / "token_type.json").read_text()))
|
||||
ids_list = [chat_ids(tok, p) for p in PROMPTS]
|
||||
args = list(cfg["args"])
|
||||
if isinstance(cfg.get("gpu"), list) and "--layer-split" not in args: # several cards: measured as it runs
|
||||
args += ["--layer-split", str(cfg.get("layer_split") or "auto")]
|
||||
return measure(args, ids_list, start_engine, say)
|
||||
return measure(engine_args(cfg), ids_list, start_engine, say)
|
||||
|
||||
|
||||
def engine_args(cfg: dict) -> list[str]:
|
||||
"""The arguments the server starts this config's engine with (serve.server.engine_args), so the tuning measures
|
||||
the engine as it runs. #447: this used to read any "gpu" list as a layer split and add --layer-split auto, which
|
||||
broke a one-card config with a helper card for the expert tier ("gpu": [0] + --expert-cache-device1: the engine
|
||||
refuses a split that leaves no card without a stage), missed the "0,2" spelling, and ignored split_skip_if_fits."""
|
||||
from serve.server import engine_args as server_args
|
||||
return server_args(cfg)
|
||||
|
||||
|
||||
def engine_error(log: str | None, since: int = 0) -> str | None:
|
||||
"""#447: the engine's own reason for a failed start - the last line of its log written after byte `since` that is
|
||||
its own ("strata ..." or "ERR ..."), else the last line there; None without a log or a new line."""
|
||||
if not log:
|
||||
return None
|
||||
try:
|
||||
with open(log, "rb") as f:
|
||||
f.seek(since)
|
||||
lines = [x.strip() for x in f.read()[-16384:].decode("utf-8", "replace").splitlines() if x.strip()]
|
||||
except OSError:
|
||||
return None
|
||||
return next((x for x in reversed(lines) if x.startswith(("strata", "ERR"))), lines[-1] if lines else None)
|
||||
|
||||
|
||||
def measure(base_args: list[str], ids_list, start_engine, say=print) -> dict:
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
tools/hip/build_windows.bat runs it after the build. The zip holds:
|
||||
|
||||
strata.exe, strata-device.exe, BUILD.json (backend "hip", the archs, the ROCm and hipBLASLt versions)
|
||||
amdhip64_7.dll, amd_comgr.dll the HIP runtime beside the exes too (#468 #461: before System32)
|
||||
rocm/bin/*.dll the ROCm DLLs the two programs load (their import tables, followed
|
||||
through the ROCm DLLs, + amd_comgr.dll, which the HIP runtime loads
|
||||
by name) and the Microsoft C++ runtime they import
|
||||
@@ -104,6 +105,12 @@ def main() -> int:
|
||||
sys.exit("VCToolsRedistDir is not set: run this from a Visual Studio x64 developer prompt (vcvars64.bat)")
|
||||
for d in sorted(crt):
|
||||
shutil.copy2(cdir / d, stage / "rocm" / "bin" / d)
|
||||
# #468 #461: the HIP runtime (and the compiler it loads by name) next to the exes as well - Windows searches the
|
||||
# exe's folder before System32, where an AMD driver may install its own amdhip64_7.dll (found before PATH's
|
||||
# rocm/bin); rocBLAS/hipBLASLt stay in rocm/bin, which they resolve their kernel libraries and ../.kpack from
|
||||
for pat in ("amdhip64_*.dll", "amd_comgr*.dll"):
|
||||
for f in (stage / "rocm" / "bin").glob(pat):
|
||||
shutil.copy2(f, stage / f.name)
|
||||
|
||||
# the GEMM kernels and the libraries' own device code, for these archs
|
||||
lib = stage / "rocm" / "bin" / "rocblas" / "library"
|
||||
|
||||
+4
-2
@@ -31,8 +31,10 @@ import urllib.request
|
||||
# repository no longer has it, the current files are read instead, with a message (resolve_repo).
|
||||
PINNED_REVISION = "de4b8e4d43b917e7706784d8bb445c9af86a3540"
|
||||
REVISION = os.environ.get("STRATA_MTP_REVISION") or PINNED_REVISION
|
||||
REPO = "https://huggingface.co/Qwen/Qwen3.8-Flash-Next/resolve/%s/" % REVISION
|
||||
PINNED = "https://huggingface.co/Qwen/Qwen3.8-Flash-Next/resolve/%s/" % PINNED_REVISION # SHA256's revision
|
||||
# #495: HF_ENDPOINT (a mirror, e.g. https://hf-mirror.com) serves the same revision; the SHA256 checks below still apply
|
||||
HF_ENDPOINT = (os.environ.get("HF_ENDPOINT") or "").strip().rstrip("/") or "https://huggingface.co"
|
||||
REPO = HF_ENDPOINT + "/Qwen/Qwen3.8-Flash-Next/resolve/%s/" % REVISION
|
||||
PINNED = HF_ENDPOINT + "/Qwen/Qwen3.8-Flash-Next/resolve/%s/" % PINNED_REVISION # SHA256's revision
|
||||
DTYPE_BYTES = {"BF16": 2, "F16": 2, "F32": 4, "F8_E4M3": 1, "I64": 8, "I32": 4}
|
||||
BAD = 3 # `verify`'s exit code: a tensor is missing or corrupt
|
||||
|
||||
|
||||
@@ -11,6 +11,7 @@ import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(ROOT / "tools"))
|
||||
@@ -91,6 +92,53 @@ class Calibrate(unittest.TestCase):
|
||||
self.assertIsNone(CAL.arg_value(b, "--pool-workers"))
|
||||
self.assertEqual(b.count("--spec-min-p"), 1)
|
||||
|
||||
def test_engine_args_are_the_servers(self):
|
||||
# #447: a "gpu" list is a layer split only with two or more cards; "0,2" is one too; split_skip_if_fits counts
|
||||
helper = {"args": BASE + ["--expert-cache-device1", "1800"], "gpu": [0], "env": {"CUDA_VISIBLE_DEVICES": "0,1"}}
|
||||
self.assertEqual(CAL.engine_args(helper), helper["args"]) # one stage + a helper card: no split
|
||||
self.assertEqual(CAL.engine_args({"args": BASE, "gpu": 1}), BASE)
|
||||
self.assertEqual(CAL.engine_args({"args": BASE, "gpu": [0, 2]}), BASE + ["--layer-split", "auto"])
|
||||
self.assertEqual(CAL.engine_args({"args": BASE, "gpu": "0,2", "layer_split": "18"}),
|
||||
BASE + ["--layer-split", "18"])
|
||||
self.assertEqual(CAL.engine_args({"args": BASE, "gpu": [0, 1], "split_skip_if_fits": True}),
|
||||
BASE + ["--layer-split", "auto", "--split-skip-if-fits"])
|
||||
own = BASE + ["--layer-split", ""] # the config's own value wins
|
||||
self.assertEqual(CAL.engine_args({"args": own, "gpu": [0, 1]}), own)
|
||||
|
||||
def test_run_measures_with_the_servers_args(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
tok = Path(d)
|
||||
(tok / "vocab.json").write_text(json.dumps({"a": 0, "b": 1}))
|
||||
(tok / "merges.txt").write_text("")
|
||||
(tok / "token_type.json").write_text(json.dumps([1, 1]))
|
||||
seen = []
|
||||
saved = CAL.measure
|
||||
CAL.measure = lambda args, ids_list, start_engine, say=print: seen.append(args) or {}
|
||||
fake = type("ST", (), {"Tokenizer": lambda *a: type("T", (), {"encode": lambda s, t, **k: [0]})()})
|
||||
try:
|
||||
with mock.patch.dict(sys.modules, {"strata_tokenizer": fake}):
|
||||
CAL.run({"tokenizer": d, "args": list(BASE), "gpu": [0]}, start_engine=lambda a: None)
|
||||
finally:
|
||||
CAL.measure = saved
|
||||
self.assertEqual(seen, [BASE])
|
||||
|
||||
def test_engine_error(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
log = Path(d) / "strata-x.log"
|
||||
log.write_text("strata generate: an old error\n", encoding="utf-8")
|
||||
since = log.stat().st_size
|
||||
with open(log, "a", encoding="utf-8") as f:
|
||||
f.write("loading ...\nstrata generate: --expert-cache-remote with a layer split needs a GPU that runs "
|
||||
"no stage (2 visible, 2 used by the split)\n\n")
|
||||
self.assertEqual(CAL.engine_error(str(log), since), "strata generate: --expert-cache-remote with a layer "
|
||||
"split needs a GPU that runs no stage (2 visible, 2 used by the split)")
|
||||
with open(log, "a", encoding="utf-8") as f:
|
||||
f.write("Segmentation fault\n")
|
||||
self.assertIn("--expert-cache-remote", CAL.engine_error(str(log), since)) # the engine's own line first
|
||||
self.assertEqual(CAL.engine_error(str(log), log.stat().st_size), None) # nothing new since
|
||||
self.assertIsNone(CAL.engine_error(None))
|
||||
self.assertIsNone(CAL.engine_error(str(Path(d) / "missing.log")))
|
||||
|
||||
def test_worker_candidates(self):
|
||||
self.assertEqual(CAL.worker_candidates(6), [6, 4, 3])
|
||||
self.assertEqual(CAL.worker_candidates(23), [23, 15, 12])
|
||||
@@ -152,6 +200,35 @@ class SetupIntegration(unittest.TestCase):
|
||||
self.assertEqual(json.loads(cfg_path.read_text())["args"], BASE)
|
||||
self.assertIsNone(self.S.saved_calibration({"args": BASE, "model_name": "m"}))
|
||||
|
||||
def test_failed_calibration_says_the_engines_reason(self):
|
||||
# #447: the engine's last error line is printed, and the --calibrate run repeats the failure before it starts
|
||||
import contextlib
|
||||
import io
|
||||
log = Path(self.tmp.name) / "strata-q2_0.log"
|
||||
log.write_text("strata generate: an error of an earlier start\n", encoding="utf-8")
|
||||
cfg_path = Path(self.tmp.name) / "strata-q2_0.json"
|
||||
cfg_path.write_text(json.dumps({"exe": "x", "args": list(BASE), "model_name": "m", "log": str(log)}))
|
||||
|
||||
def boom(*a, **k):
|
||||
with open(log, "a", encoding="utf-8") as f:
|
||||
f.write("strata generate: --expert-cache-remote with a layer split needs a GPU that runs no stage\n")
|
||||
raise RuntimeError(f"the engine exited before it was ready (see {log})")
|
||||
CAL.run = boom
|
||||
out = io.StringIO()
|
||||
with contextlib.redirect_stdout(out):
|
||||
self.assertFalse(self.S.calibrate_config(cfg_path))
|
||||
self.assertIn("the engine said: strata generate: --expert-cache-remote with a layer split", out.getvalue())
|
||||
self.assertNotIn("an earlier start", out.getvalue())
|
||||
out = io.StringIO()
|
||||
with mock.patch.object(self.S, "ROOT", Path(self.tmp.name)), \
|
||||
mock.patch.object(self.S, "data_folder", lambda d: (Path(self.tmp.name) / "data", [])), \
|
||||
mock.patch.object(self.S, "update_installed_engine", lambda *a: None), \
|
||||
mock.patch.object(sys, "argv", ["setup.py", "--calibrate", "--no-start"]), \
|
||||
contextlib.redirect_stdout(out):
|
||||
self.assertEqual(self.S.main(), 0)
|
||||
self.assertIn("this PC is NOT tuned: the tuning failed (the reason is above); the model keeps the default "
|
||||
"settings", out.getvalue())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -117,6 +117,40 @@ class GpuLists(unittest.TestCase):
|
||||
else:
|
||||
setup.os.environ["ROCM_PATH"] = old
|
||||
|
||||
def test_runtime_only_system_rocm_falls_back_to_the_wheels(self):
|
||||
"""#446: a system ROCm 7 with hipcc and libhipblas but no HIP development files (no hip-lang CMake package, no
|
||||
hip_runtime.h) is not used for the build: a warning says what is missing and the wheels path follows (here
|
||||
it stops at the two-family check, which only the wheels path makes); with the files it is used as before."""
|
||||
said = []
|
||||
setup.say = lambda msg="": said.append(msg) # tearDown puts the real one back
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
sysroot = Path(d)
|
||||
for rel, text in (("bin/hipcc", ""), ("lib/libhipblas.so.3", ""),
|
||||
("include/rocm-core/rocm_version.h",
|
||||
"#define ROCM_VERSION_MAJOR 7\n#define ROCM_VERSION_MINOR 14\n")):
|
||||
(sysroot / rel).parent.mkdir(parents=True, exist_ok=True)
|
||||
(sysroot / rel).write_text(text)
|
||||
with mock.patch.dict(setup.os.environ, {"ROCM_PATH": d}):
|
||||
with self.assertRaises(SystemExit):
|
||||
setup.rocm_root(["gfx1100", "gfx1201"])
|
||||
self.assertIn(f"the ROCm in {sysroot} has no HIP development files (lib/cmake/hip-lang/hip-lang-"
|
||||
"config.cmake, include/hip/hip_runtime.h): using AMD's wheels", "\n".join(said))
|
||||
self.assertIn("two GPU families", "\n".join(said))
|
||||
for lib in ("lib64", "lib"): # either place CMake looks
|
||||
with self.subTest(lib=lib):
|
||||
cfg = sysroot / lib / "cmake/hip-lang/hip-lang-config.cmake"
|
||||
cfg.parent.mkdir(parents=True, exist_ok=True)
|
||||
cfg.write_text("")
|
||||
said.clear()
|
||||
with self.assertRaises(SystemExit):
|
||||
setup.rocm_root(["gfx1100", "gfx1201"])
|
||||
self.assertIn("(include/hip/hip_runtime.h)", "\n".join(said))
|
||||
(sysroot / "include/hip").mkdir(parents=True, exist_ok=True)
|
||||
(sysroot / "include/hip/hip_runtime.h").write_text("")
|
||||
self.assertEqual(setup.rocm_root(["gfx1100", "gfx1201"]), (sysroot, [str(sysroot / "lib")]))
|
||||
cfg.unlink()
|
||||
(sysroot / "include/hip/hip_runtime.h").unlink()
|
||||
|
||||
def test_build_for_every_arch(self):
|
||||
"""build_engine_hip compiles for the set of the chosen cards' archs and records it in BUILD.json."""
|
||||
calls = {}
|
||||
@@ -311,5 +345,32 @@ class WindowsHipVision(unittest.TestCase):
|
||||
self.assertEqual(setup.hip_vision("cpu"), "cpu")
|
||||
|
||||
|
||||
class HipRuntimeBesideExe(unittest.TestCase):
|
||||
"""#468 #461: the bundled HIP runtime (and amd_comgr) goes next to strata.exe, so an AMD driver's System32 copy is
|
||||
not found first; rocBLAS and the rest stay in rocm/bin."""
|
||||
|
||||
def test_copies_only_the_runtime(self):
|
||||
import json
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
eng = Path(d)
|
||||
rb = eng / "rocm" / "bin"
|
||||
rb.mkdir(parents=True)
|
||||
(eng / "BUILD.json").write_text(json.dumps({"backend": "hip", "lib_dirs": ["rocm/bin"]}))
|
||||
for n, data in (("amdhip64_7.dll", b"hip-3686"), ("amd_comgr.dll", b"comgr"), ("rocblas.dll", b"blas")):
|
||||
(rb / n).write_bytes(data)
|
||||
setup.hip_runtime_beside_exe(eng)
|
||||
self.assertEqual((eng / "amdhip64_7.dll").read_bytes(), b"hip-3686")
|
||||
self.assertEqual((eng / "amd_comgr.dll").read_bytes(), b"comgr")
|
||||
self.assertFalse((eng / "rocblas.dll").exists()) # finds its kernels relative to rocm/bin
|
||||
(rb / "amdhip64_7.dll").write_bytes(b"hip-3690-newer") # a newer zip: replaced on the next start
|
||||
setup.hip_runtime_beside_exe(eng)
|
||||
self.assertEqual((eng / "amdhip64_7.dll").read_bytes(), b"hip-3690-newer")
|
||||
|
||||
def test_no_rocm_bin_is_a_no_op(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
setup.hip_runtime_beside_exe(Path(d)) # no BUILD.json (a CUDA or Linux engine)
|
||||
self.assertEqual(list(Path(d).iterdir()), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -92,6 +92,64 @@ class GgufDirShards(unittest.TestCase):
|
||||
self.assertEqual(self.shards(["a-00001-of-00002.gguf", "b-00001-of-00002.gguf"]), want) # ambiguous
|
||||
|
||||
|
||||
class GgufDirUnsupported(unittest.TestCase):
|
||||
"""#444: --gguf-dir with GGUFs Strata cannot run (Unsloth's UD-IQ3_XXS, UD-Q2_K_XL) says so, names the files it
|
||||
runs (ISTA-DASLab's GSQ-RCO, the Coder's, Unsloth's UD-Q4_K_XL) and the --family/--model of the usable ones."""
|
||||
UD = ["Qwen3.8-Flash-Next-UD-IQ3_XXS-%05d-of-00003.gguf" % i for i in range(1, 4)] + \
|
||||
["Qwen3.8-Flash-Next-UD-Q2_K_XL-%05d-of-00003.gguf" % i for i in range(1, 4)]
|
||||
|
||||
def problem(self, names, family="qwen", model="IQ3_XXS"):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
for n in names:
|
||||
(Path(d) / n).write_bytes(b"")
|
||||
fam = setup.FAMILIES[family]
|
||||
first = setup.gguf_dir_shards(Path(d), fam, model)[0]
|
||||
p = setup.gguf_dir_problem(Path(d), first, fam, model)
|
||||
return p and (p[0].replace(d, "<D>"), p[1])
|
||||
|
||||
def test_quant_names(self):
|
||||
for name, want in (("Qwen3.8-Flash-Next-UD-IQ3_XXS-00001-of-00003.gguf", "UD-IQ3_XXS"),
|
||||
("Qwen3.8-Flash-Next-UD-Q2_K_XL-00001-of-00003.gguf", "UD-Q2_K_XL"),
|
||||
("model-Q4_K_M.gguf", "Q4_K_M"),
|
||||
("Qwen3.8-Flash-Next-GSQ-RCO-IQ3_XXS-00001-of-00002.gguf", None),
|
||||
("Qwen3.8-Flash-Next-UD-Q4_K_XL-00001-of-00004.gguf", None),
|
||||
("my-IQ3_XXS-00001-of-00003.gguf", None), ("mmproj-F16.gguf", None),
|
||||
("a-00001-of-00002.gguf", None)):
|
||||
self.assertEqual(setup.gguf_unsupported(name), want, name)
|
||||
|
||||
def test_an_unsloth_ud_file_is_not_taken_for_a_gsq_rco_size(self):
|
||||
msg, hint = self.problem(self.UD) # IQ3_XXS: UD-IQ3_XXS has the size in its name
|
||||
self.assertEqual(msg, "Qwen3.8-Flash-Next-UD-IQ3_XXS-00001-of-00003.gguf is UD-IQ3_XXS, a GGUF Strata "
|
||||
"cannot run")
|
||||
self.assertIn("ISTA-DASLab's GSQ-RCO files", hint)
|
||||
self.assertIn("Unsloth's UD-Q4_K_XL only", hint)
|
||||
|
||||
def test_a_folder_without_the_choice_names_what_is_there(self):
|
||||
gsq = ["Qwen3.8-Flash-Next-GSQ-RCO-IQ1_M-%05d-of-00002.gguf" % i for i in (1, 2)]
|
||||
msg, hint = self.problem(self.UD + gsq, "unsloth", "UD-Q4_K_XL")
|
||||
self.assertEqual(msg, "<D> has no Qwen3.8-Flash-Next (Unsloth) UD-Q4_K_XL file")
|
||||
self.assertIn("Not usable here: Qwen3.8-Flash-Next-UD-IQ3_XXS-00001-of-00003.gguf, "
|
||||
"Qwen3.8-Flash-Next-UD-Q2_K_XL-00001-of-00003.gguf", hint)
|
||||
self.assertIn("Usable here: --family coder --model IQ1_M", hint)
|
||||
|
||||
def test_usable_folders_are_left_alone(self):
|
||||
good = ["Qwen3.8-Flash-Next-GSQ-RCO-IQ3_XXS-%05d-of-00002.gguf" % i for i in (1, 2)]
|
||||
self.assertIsNone(self.problem(good + self.UD))
|
||||
self.assertIsNone(self.problem(["my-IQ3_XXS-%05d-of-00003.gguf" % i for i in (1, 2, 3)]))
|
||||
self.assertIsNone(self.problem([])) # nothing there: check_shards' "missing", as before
|
||||
|
||||
def test_the_size_of_another_family(self):
|
||||
from test_setup_golden import PROFILES, install
|
||||
ram, found = PROFILES["64GB-1x32GB"]
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
code, out, _, _ = install(ram, found, ["--gguf-dir", d, "--family", "unsloth", "--model", "IQ3_XXS"])
|
||||
self.assertEqual(code, 1)
|
||||
self.assertIn("has no IQ3_XXS model file", out)
|
||||
self.assertIn("choose one of: UD-Q4_K_XL (or IQ3_XXS: --family qwen --model IQ3_XXS, --family swift --model "
|
||||
"IQ3_XXS)", out)
|
||||
self.assertIn("Strata runs ISTA-DASLab's GSQ-RCO files", out)
|
||||
|
||||
|
||||
class ExperimentalSm60(unittest.TestCase):
|
||||
"""#295: Pascal (6.x) and Volta (7.0) only with STRATA_EXPERIMENTAL_SM60=1, built with -DSTRATA_EXPERIMENTAL_SM60=ON
|
||||
and a CUDA 12.x toolkit; nothing changes without the variable."""
|
||||
@@ -209,5 +267,68 @@ class HipVision(unittest.TestCase):
|
||||
self.assertNotIn("vision_src", meta)
|
||||
|
||||
|
||||
class VramReserve(unittest.TestCase):
|
||||
"""#493: --vram-reserve-mib N - VRAM the engine leaves free for other programs - goes into the config only when
|
||||
given (the default config is tools/test_setup_golden.py's, unchanged), at setup and on a start."""
|
||||
|
||||
def setUp(self):
|
||||
from test_setup_golden import PROFILES
|
||||
self.ram, self.found = PROFILES["64GB-1x32GB"]
|
||||
|
||||
def install(self, *extra):
|
||||
from test_setup_golden import install
|
||||
return install(self.ram, self.found, ["--family", "qwen", "--model", "Q2_0", "--no-start", *extra])
|
||||
|
||||
def test_given_at_setup(self):
|
||||
code, out, cfg, _ = self.install("--vram-reserve-mib", "2048")
|
||||
self.assertEqual(code, 0, out)
|
||||
a = cfg["args"]
|
||||
self.assertEqual(a.count("--vram-reserve-mib"), 1)
|
||||
self.assertEqual(a[a.index("--vram-reserve-mib") + 1], "2048")
|
||||
self.assertIn("VRAM kept free for other programs: 2048 MiB", out)
|
||||
code, out, cfg, _ = self.install()
|
||||
self.assertNotIn("--vram-reserve-mib", cfg["args"]) # not given: the engine's default, as before
|
||||
|
||||
def test_with_images_one_value(self):
|
||||
code, out, cfg, _ = self.install("--vision", "cpu", "--vram-reserve-mib", "1500")
|
||||
self.assertEqual(code, 0, out)
|
||||
a = cfg["args"]
|
||||
self.assertEqual(a.count("--vram-reserve-mib"), 1)
|
||||
self.assertEqual(a[a.index("--vram-reserve-mib") + 1], "1500")
|
||||
|
||||
def test_a_negative_value_is_refused(self):
|
||||
with contextlib.redirect_stderr(io.StringIO()) as err:
|
||||
code, out, cfg, _ = self.install("--vram-reserve-mib", "-1")
|
||||
self.assertEqual(code, 2)
|
||||
self.assertIn("--vram-reserve-mib takes a number of MiB", err.getvalue())
|
||||
|
||||
def test_kept_from_an_earlier_install(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
p = Path(d) / "strata-q2_0.json"
|
||||
p.write_text(json.dumps({"args": ["--max-context", "8192", "--vram-reserve-mib", "2048"]}))
|
||||
self.assertEqual(setup.choices_from_config(p)["vram_reserve_mib"], 2048)
|
||||
p.write_text(json.dumps({"args": ["--vision", "--vram-reserve-mib", "700"], "vision": {"gpu": True}}))
|
||||
self.assertIsNone(setup.choices_from_config(p)["vram_reserve_mib"]) # the images' own reserve
|
||||
p.write_text(json.dumps({"args": ["--max-context", "8192"]}))
|
||||
self.assertIsNone(setup.choices_from_config(p)["vram_reserve_mib"])
|
||||
|
||||
def test_given_on_a_start(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
exe = Path(d) / "strata.exe"
|
||||
exe.write_bytes(b"")
|
||||
p = Path(d) / "strata-q2_0.json"
|
||||
p.write_text(json.dumps({"exe": str(exe), "args": ["--kv", "int8"], "gpu": 0, "gpus_asked": True}))
|
||||
call = mock.Mock(return_value=0)
|
||||
with mock.patch.object(setup, "gpus", lambda: self.found), \
|
||||
mock.patch.object(setup, "is_wsl", lambda: False), \
|
||||
mock.patch.object(setup, "ensure_engine_for", lambda cards, path, cfg, yes: cfg), \
|
||||
mock.patch.object(setup.subprocess, "call", call):
|
||||
for n in ("2048", "1024"):
|
||||
_, out = quiet(setup.start, p, None, None, False, True, None, {"vram_reserve_mib": int(n)})
|
||||
self.assertEqual(json.loads(p.read_text())["args"], ["--kv", "int8", "--vram-reserve-mib", n])
|
||||
self.assertIn("1024 MiB of VRAM kept free", out)
|
||||
self.assertTrue(call.called)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -49,5 +49,29 @@ class Refresh(unittest.TestCase):
|
||||
self.assertEqual(hashlib.sha256((rt / "draft_vocab.bin").read_bytes()).hexdigest(), want)
|
||||
|
||||
|
||||
class SmallCardNote(unittest.TestCase):
|
||||
"""#474: a card under 14 GB is told about a smaller draft subset - a note only, and not when one was chosen."""
|
||||
|
||||
def test_note_on_a_small_card_only(self):
|
||||
note = setup.draft_vocab_note(12.0, None)
|
||||
self.assertTrue(note)
|
||||
self.assertIn("--draft-vocab en", " ".join(note))
|
||||
self.assertIn("the draft head does not fit", " ".join(note))
|
||||
for vram, chosen in ((16.0, None), (24.0, None), (12.0, "en"), (12.0, "cyrillic"), (12.0, "cjk"), (0.0, None)):
|
||||
self.assertEqual(setup.draft_vocab_note(vram, chosen), [], (vram, chosen))
|
||||
|
||||
def test_sizes_follow_the_shipped_subsets(self):
|
||||
# the MiB the note gives scale with the subsets' token counts (4 bytes per token id in data/)
|
||||
sizes = {}
|
||||
for choice, name in setup.DRAFT_VOCABS.items():
|
||||
p = ROOT / "data" / name
|
||||
if not p.exists():
|
||||
self.skipTest(f"data/{name} is not in this checkout")
|
||||
sizes[choice] = p.stat().st_size // 4
|
||||
for choice in ("en", "cyrillic"):
|
||||
want = setup.DRAFT_VOCAB_MIB["cjk"] * sizes[choice] / sizes["cjk"]
|
||||
self.assertAlmostEqual(setup.DRAFT_VOCAB_MIB[choice], want, delta=3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -98,6 +98,48 @@ class HuggingFacePins(unittest.TestCase):
|
||||
finally:
|
||||
mtp_fetch.REPO = pinned
|
||||
|
||||
def test_hf_endpoint(self):
|
||||
# #495: HF_ENDPOINT (a mirror) serves the same pinned revision; a trailing slash and blanks are dropped
|
||||
repo = "ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF"
|
||||
with mock.patch.dict(setup.os.environ, {}, clear=False):
|
||||
setup.os.environ.pop("HF_ENDPOINT", None)
|
||||
self.assertEqual(setup.hf(repo), f"https://huggingface.co/{repo}/resolve/{setup.HF_REVISIONS[repo]}/")
|
||||
for value in ("https://hf-mirror.com", "https://hf-mirror.com/", " https://hf-mirror.com/ "):
|
||||
setup.os.environ["HF_ENDPOINT"] = value
|
||||
url = setup.hf(repo)
|
||||
self.assertEqual(url, f"https://hf-mirror.com/{repo}/resolve/{setup.HF_REVISIONS[repo]}/")
|
||||
self.assertRegex(url, SHA)
|
||||
self.assertEqual(setup.hf_unpinned(url + "x.gguf"), f"https://hf-mirror.com/{repo}/resolve/main/x.gguf")
|
||||
setup.os.environ["HF_ENDPOINT"] = ""
|
||||
self.assertEqual(setup.hf_endpoint(), "https://huggingface.co")
|
||||
|
||||
def test_mtp_fetch_honours_hf_endpoint(self):
|
||||
import importlib
|
||||
import mtp_fetch
|
||||
try:
|
||||
with mock.patch.dict(mtp_fetch.os.environ, {"HF_ENDPOINT": "https://hf-mirror.com/"}):
|
||||
importlib.reload(mtp_fetch)
|
||||
self.assertTrue(mtp_fetch.REPO.startswith("https://hf-mirror.com/Qwen/Qwen3.8-Flash-Next/resolve/"))
|
||||
self.assertRegex(mtp_fetch.REPO, SHA)
|
||||
self.assertTrue(mtp_fetch.pinned()) # the tensors' SHA-256 checks still apply
|
||||
finally:
|
||||
importlib.reload(mtp_fetch)
|
||||
self.assertTrue(mtp_fetch.REPO.startswith("https://huggingface.co/"))
|
||||
|
||||
def test_step5_names_the_folder(self):
|
||||
# #495: where setup expects the model files, and how to give it files downloaded by hand
|
||||
from test_setup_golden import PROFILES, install
|
||||
ram, found = PROFILES["96GB-1x16GB"]
|
||||
with mock.patch.dict(setup.os.environ, {"HF_ENDPOINT": "https://hf-mirror.com"}):
|
||||
code, out, cfg, _ = install(ram, found, ["--family", "qwen", "--model", "IQ3_XXS", "--no-start"])
|
||||
self.assertEqual(code, 0)
|
||||
text = out.replace("\\", "/")
|
||||
self.assertRegex(text, r"The model files go in .*/models/IQ3_XXS")
|
||||
self.assertIn("put them here with their original names (Qwen3.8-Flash-Next-GSQ-RCO-IQ3_XXS-00001-of-00002.gguf",
|
||||
text)
|
||||
self.assertIn("--gguf-dir", text)
|
||||
self.assertIn("Downloading from https://hf-mirror.com (HF_ENDPOINT)", text)
|
||||
|
||||
|
||||
class Engine(unittest.TestCase):
|
||||
def setUp(self):
|
||||
|
||||
@@ -10,6 +10,7 @@ from __future__ import annotations
|
||||
import contextlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
@@ -135,6 +136,53 @@ class Context(unittest.TestCase):
|
||||
self.assertTrue(started.called)
|
||||
|
||||
|
||||
class BrokenEarlierConfig(unittest.TestCase):
|
||||
"""#459: a new copy of Strata set up like an earlier install skips an earlier config that does not parse (an
|
||||
empty strata-*.json crashed START-HERE with JSONDecodeError) and goes on as a fresh install; configs are written
|
||||
whole (a temporary file moved over the old one)."""
|
||||
RAM64, GPU32 = PROFILES["64GB-1x32GB"]
|
||||
|
||||
def setup_with(self, files):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
for i, (name, text) in enumerate(files): # in order, oldest first
|
||||
(Path(d) / name).write_text(text, encoding="utf-8")
|
||||
os.utime(Path(d) / name, (1_700_000_000 + i, 1_700_000_000 + i))
|
||||
return install(self.RAM64, self.GPU32, [], extra=[
|
||||
mock.patch.object(setup, "other_installs", lambda settings: [Path(d)]),
|
||||
mock.patch.object(setup, "start", mock.Mock(return_value=0))])
|
||||
|
||||
def test_an_empty_config_is_skipped(self):
|
||||
for text, why in (("", "the file is empty"), ("{\"args\": [", "not valid JSON"), ("[1]", "not a JSON object")):
|
||||
with self.subTest(text=text):
|
||||
code, out, cfg, _ = self.setup_with([("strata-iq3_s.json", text)])
|
||||
self.assertEqual(code, 0, out)
|
||||
self.assertIn("skipped the earlier config ", out)
|
||||
self.assertIn(f"strata-iq3_s.json ({why}): setting this copy up without it", out)
|
||||
self.assertNotIn("Found your earlier install", out)
|
||||
self.assertIsNotNone(cfg) # the fresh install's config
|
||||
|
||||
def test_the_newest_readable_config_is_used(self):
|
||||
good = json.dumps({"args": ["--max-context", "262144", "--kv", "int8"], "port": 8080})
|
||||
code, out, cfg, _ = self.setup_with([("strata-iq3_s.json", good), ("strata-q2_0.json", "")])
|
||||
self.assertEqual(code, 0, out)
|
||||
self.assertIn("strata-q2_0.json (the file is empty)", out)
|
||||
self.assertIn("Found your earlier install", out)
|
||||
self.assertIn("(iq3_s)", out)
|
||||
self.assertEqual(arg(cfg, "--max-context"), "262144")
|
||||
|
||||
def test_write_config_leaves_no_partial_file(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
p = Path(d) / "strata-iq3_s.json"
|
||||
p.write_text("old", encoding="utf-8")
|
||||
setup.write_config(p, {"args": ["--kv", "int8"]})
|
||||
self.assertEqual(p.read_text(encoding="utf-8"), json.dumps({"args": ["--kv", "int8"]}, indent=1))
|
||||
self.assertEqual([f.name for f in Path(d).iterdir()], ["strata-iq3_s.json"]) # no .tmp left
|
||||
with mock.patch.object(Path, "write_text", side_effect=OSError(28, "No space left on device")):
|
||||
with self.assertRaises(OSError):
|
||||
setup.write_config(p, {"args": []})
|
||||
self.assertEqual(json.loads(p.read_text(encoding="utf-8")), {"args": ["--kv", "int8"]}) # kept whole
|
||||
|
||||
|
||||
class LowRamGpus(unittest.TestCase):
|
||||
"""S2/S3 (#364 #384): one GPU recommended in the low-RAM mode, all of them when asked for."""
|
||||
|
||||
@@ -277,6 +325,42 @@ class SmallCard(unittest.TestCase):
|
||||
self.assertEqual(cfg["gpu"], 0) # one card: as before
|
||||
|
||||
|
||||
class SplitShortCard(unittest.TestCase):
|
||||
"""#448: a second card too small to lend a split's prompt chunk: the first card alone recommended, the pair kept
|
||||
when named."""
|
||||
FOUND = [card(0, "NVIDIA RTX PRO 4500 Blackwell", 31.8, "120"), card(1, "NVIDIA GeForce RTX 3080", 10.0, "86")]
|
||||
|
||||
def test_rule(self):
|
||||
self.assertEqual([g["index"] for g in setup.split_short(self.FOUND)], [1])
|
||||
self.assertEqual(setup.split_short(PROFILES["32GB-2x24GB"][1]), [])
|
||||
self.assertEqual(setup.split_short(PROFILES["47GB-2x16GB"][1]), [])
|
||||
eights = [card(i, "NVIDIA GeForce RTX 3070", 8.0, "86") for i in range(2)]
|
||||
self.assertEqual(setup.split_short(eights), []) # neither card alone reads more
|
||||
|
||||
def test_yes_takes_the_first_card(self):
|
||||
code, out, cfg, _ = install(127.8, self.FOUND, ["--family", "qwen", "--model", "Q2_0", "--no-start"])
|
||||
self.assertEqual(code, 0, out)
|
||||
self.assertEqual(cfg["gpu"], 0)
|
||||
self.assertIn("too small to lend a split its prompt buffers", out)
|
||||
|
||||
def test_named_pair_is_kept(self):
|
||||
code, out, cfg, _ = install(127.8, self.FOUND, ["--family", "qwen", "--model", "Q2_0", "--no-start",
|
||||
"--gpus", "0,1"])
|
||||
self.assertEqual(code, 0, out)
|
||||
self.assertEqual(cfg["gpu"], [0, 1])
|
||||
|
||||
def test_offer_together_defaults_to_one(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
p = Path(d) / "strata-q2_0.json"
|
||||
with mock.patch.object(setup, "gpus", lambda: self.FOUND), \
|
||||
mock.patch.object(setup, "engine_runs_on", lambda g: True):
|
||||
code, out, asked = run(setup.offer_together, p, {"args": ["--mmap-experts"]}, True)
|
||||
cfg = json.loads(p.read_text())
|
||||
self.assertIsNone(code, out)
|
||||
self.assertNotIn("gpu", cfg)
|
||||
self.assertIn("#448", out)
|
||||
|
||||
|
||||
class RamFloor(unittest.TestCase):
|
||||
"""S7: less RAM than the smallest model needs: a stop with --yes alone, a risk with --model."""
|
||||
FOUND = [card(0, "NVIDIA GeForce RTX 3060", 8.0, "86")]
|
||||
|
||||
@@ -319,6 +319,117 @@ class Main(unittest.TestCase):
|
||||
self.assertNotIn("--mmap-experts", cfg["args"])
|
||||
self.assertFalse(any("--experts-bin" in r for r in self.runs))
|
||||
|
||||
def test_split_when_the_ram_holds_it(self):
|
||||
# #498: 165 GiB and --gpus 0,1: the split without the RAM budget (the engine refuses the pair)
|
||||
code, out, cfg = self.main(["--context", "8192", "--gpus", "0,1"], n_gpus=2, ram=165.0)
|
||||
self.assertEqual(code, 0, out)
|
||||
self.assertEqual(cfg["gpu"], [0, 1])
|
||||
self.assertEqual(cfg["layer_split"], "auto")
|
||||
for flag in ("--resident-budget-gib", "--mmap-experts", "--resident-experts"):
|
||||
self.assertNotIn(flag, cfg["args"])
|
||||
self.assertIn("as you chose (--gpus): no RAM budget", out)
|
||||
self.assertIn("OS file cache", out)
|
||||
self.assertNotIn("RAM budget: ", out)
|
||||
|
||||
def test_split_refused_without_the_ram(self):
|
||||
need = setup.unsloth_split_need_gb()
|
||||
self.assertAlmostEqual(need, setup.MODELS[M]["download_gb"] + setup.UNSLOTH_RAM_LEFT_GB)
|
||||
code, out, cfg = self.main(["--context", "8192", "--gpus", "0,1"], n_gpus=2, ram=127.8)
|
||||
self.assertEqual(code, 0, out)
|
||||
self.assertIn("runs on one GPU here", out)
|
||||
self.assertNotIn("layer_split", cfg)
|
||||
self.assertEqual(cfg["args"][cfg["args"].index("--resident-budget-gib") + 1], "71") # all of them
|
||||
|
||||
def test_an_explicit_budget_keeps_one_gpu(self):
|
||||
code, out, cfg = self.main(["--context", "8192", "--gpus", "0,1", "--resident-budget-gib", "60"], n_gpus=2,
|
||||
ram=165.0)
|
||||
self.assertEqual(code, 0, out)
|
||||
self.assertIn("--resident-budget-gib has no layer split", out)
|
||||
self.assertNotIn("layer_split", cfg)
|
||||
self.assertEqual(cfg["args"][cfg["args"].index("--resident-budget-gib") + 1], "60")
|
||||
|
||||
def test_yes_alone_keeps_one_gpu(self):
|
||||
code, out, cfg = self.main(["--context", "8192"], n_gpus=2, ram=165.0)
|
||||
self.assertEqual(code, 0, out)
|
||||
self.assertIn("runs on one GPU: using", out)
|
||||
self.assertIn("--gpus 0,1 shares it", out)
|
||||
self.assertNotIn("layer_split", cfg)
|
||||
self.assertIn("--resident-budget-gib", cfg["args"])
|
||||
|
||||
|
||||
class LayerSplit(unittest.TestCase):
|
||||
"""#498: UD-Q4_K_XL across GPUs - asked at setup (one GPU by default), and a start with --gpus no longer keeps
|
||||
the RAM budget the engine refuses with a split (it exited with code 2)."""
|
||||
|
||||
def setUp(self):
|
||||
sys.path.insert(0, str(ROOT / "tools"))
|
||||
from test_setup_golden import card
|
||||
self.found = [card(i, "NVIDIA GeForce RTX 3090", 24.0, "86") for i in range(2)]
|
||||
|
||||
def test_asked_one_gpu_by_default(self):
|
||||
from test_setup_golden import install
|
||||
argv = ["--family", "unsloth", "--model", M, "--context", "32768", "--no-start"]
|
||||
code, out, cfg, asked = install(165.0, self.found, argv, answers={"which GPUs?": "1", "Which GPUs?": "1"})
|
||||
self.assertEqual(code, 0, out)
|
||||
self.assertTrue(any(f"{M}: which GPUs?" in q and "[1]" in q for q in asked), asked)
|
||||
self.assertEqual(cfg["gpu"], 0)
|
||||
self.assertIn("--resident-budget-gib", cfg["args"])
|
||||
code, out, cfg, asked = install(165.0, self.found, argv, answers={f"{M}: which GPUs?": "2"})
|
||||
self.assertEqual(code, 0, out)
|
||||
self.assertEqual(cfg["gpu"], [0, 1])
|
||||
self.assertEqual(cfg["layer_split"], "auto")
|
||||
self.assertNotIn("--resident-budget-gib", cfg["args"])
|
||||
|
||||
def start(self, ram, gpu=(0, 1), offered=None):
|
||||
from test_setup_risk import run
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
exe = Path(d) / "strata.exe"
|
||||
exe.write_bytes(b"")
|
||||
p = Path(d) / "strata-unsloth-ud-q4_k_xl.json"
|
||||
before = {"exe": str(exe), "args": ["--pack", "p", "--resident-budget-gib", "40", "--kv", "int8"],
|
||||
"gpu": 0, "gpus_asked": offered is None}
|
||||
p.write_text(json.dumps(before))
|
||||
call = mock.Mock(return_value=0)
|
||||
with mock.patch.object(setup, "gpus", lambda: self.found), \
|
||||
mock.patch.object(setup, "ram_gb", lambda: ram), \
|
||||
mock.patch.object(setup, "engine_runs_on", lambda g: True), \
|
||||
mock.patch.object(setup, "ensure_engine_for", lambda cards, path, cfg, yes: cfg), \
|
||||
mock.patch.object(setup, "is_wsl", lambda: False), \
|
||||
mock.patch.object(setup.subprocess, "call", call):
|
||||
code, out, asked = run(setup.start, p, None, list(gpu) if gpu else None, False, offered is None,
|
||||
stdin=offered)
|
||||
return code, out, json.loads(p.read_text()), call.called
|
||||
|
||||
def test_start_with_gpus_drops_the_budget(self):
|
||||
code, out, cfg, started = self.start(165.0)
|
||||
self.assertIsNone(code, out)
|
||||
self.assertTrue(started)
|
||||
self.assertEqual(cfg["gpu"], [0, 1])
|
||||
self.assertEqual(cfg["args"], ["--pack", "p", "--kv", "int8"])
|
||||
self.assertIn("no RAM budget", out)
|
||||
|
||||
def test_start_with_gpus_refused_without_the_ram(self):
|
||||
code, out, cfg, started = self.start(63.7)
|
||||
self.assertEqual(code, 1)
|
||||
self.assertFalse(started)
|
||||
self.assertIn("cannot share its RAM budget across GPUs", out)
|
||||
self.assertIn("--gpu N", out)
|
||||
self.assertEqual(cfg["gpu"], 0) # nothing saved
|
||||
self.assertIn("--resident-budget-gib", cfg["args"])
|
||||
|
||||
def test_offered_at_a_start(self):
|
||||
code, out, cfg, _ = self.start(165.0, gpu=None, offered="") # Enter: one GPU (the default here)
|
||||
self.assertIsNone(code, out)
|
||||
self.assertEqual(cfg["gpu"], 0)
|
||||
self.assertIn("--resident-budget-gib", cfg["args"])
|
||||
self.assertIn("about twice as fast in #498", out)
|
||||
code, out, cfg, _ = self.start(165.0, gpu=None, offered="y")
|
||||
self.assertEqual(cfg["gpu"], [0, 1])
|
||||
self.assertNotIn("--resident-budget-gib", cfg["args"])
|
||||
code, out, cfg, _ = self.start(63.7, gpu=None, offered="y") # not offered: the RAM does not hold it
|
||||
self.assertEqual(cfg["gpu"], 0)
|
||||
self.assertNotIn("Use both", out)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,87 @@
|
||||
"""Tests for `setup.py --update` (#475, UPDATE.bat / update.sh): it refreshes what a start would - the Python packages,
|
||||
the engine, each model's config and draft subset - and never starts the model. Every outside effect is mocked: no
|
||||
GPU, no downloads, nothing written outside a temp folder.
|
||||
|
||||
python -m unittest tools.test_setup_update
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(ROOT))
|
||||
import setup # noqa: E402
|
||||
|
||||
|
||||
class Update(unittest.TestCase):
|
||||
def config(self, d: Path, **extra) -> Path:
|
||||
rt = d / "mtp" / "rt"
|
||||
rt.mkdir(parents=True)
|
||||
p = d / "strata-iq3_s.json"
|
||||
p.write_text(json.dumps({"model_name": "Qwen IQ3_S", "exe": str(d / "engine" / setup.EXE),
|
||||
"args": ["--native", "x", "--mtp", str(rt)], **extra}), encoding="utf-8")
|
||||
return p
|
||||
|
||||
def run_update(self, have, argv=()):
|
||||
a = mock.Mock(**{"build": False, "prebuilt": "URL", **dict(argv)})
|
||||
out = io.StringIO()
|
||||
with mock.patch.object(setup, "pip_install") as pip, \
|
||||
mock.patch.object(setup, "update_installed_engine") as eng, \
|
||||
mock.patch.object(setup, "refresh_draft_vocab") as dv, \
|
||||
mock.patch.object(setup, "start") as start, \
|
||||
mock.patch("subprocess.call") as call, \
|
||||
contextlib.redirect_stdout(out):
|
||||
rc = setup.update_install(have, a)
|
||||
return rc, pip, eng, dv, start, call, out.getvalue()
|
||||
|
||||
def test_refreshes_without_starting(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
p = self.config(Path(d), draft_vocab="en")
|
||||
rc, pip, eng, dv, start, call, out = self.run_update([p])
|
||||
self.assertEqual(rc, 0)
|
||||
pip.assert_called_once()
|
||||
eng.assert_called_once_with("URL")
|
||||
dv.assert_called_once()
|
||||
self.assertEqual(dv.call_args[0][1], "en") # the model's own subset is kept
|
||||
start.assert_not_called()
|
||||
call.assert_not_called()
|
||||
self.assertIn("Strata is updated", out)
|
||||
|
||||
def test_build_keeps_the_compiled_engine_path(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
p = self.config(Path(d))
|
||||
rc, pip, eng, dv, *_ = self.run_update([p], {"build": True})
|
||||
self.assertEqual(rc, 0)
|
||||
eng.assert_not_called()
|
||||
self.assertEqual(dv.call_args[0][1], "cjk")
|
||||
|
||||
def test_nothing_installed(self):
|
||||
rc, pip, eng, dv, start, call, out = self.run_update([])
|
||||
self.assertEqual(rc, 0)
|
||||
for m in (pip, eng, dv, start, call):
|
||||
m.assert_not_called()
|
||||
self.assertIn("START-HERE.bat", out)
|
||||
|
||||
def test_main_update_never_starts(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
p = self.config(Path(d))
|
||||
with mock.patch.object(sys, "argv", ["setup.py", "--update"]), \
|
||||
mock.patch.object(setup, "data_folder", return_value=(Path(d), [])), \
|
||||
mock.patch.object(setup, "installed_configs", return_value=[p]), \
|
||||
mock.patch.object(setup, "update_install", return_value=0) as up, \
|
||||
mock.patch.object(setup, "start") as start, \
|
||||
contextlib.redirect_stdout(io.StringIO()):
|
||||
self.assertEqual(setup.main(), 0)
|
||||
up.assert_called_once()
|
||||
start.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,31 @@
|
||||
#!/bin/sh
|
||||
# Update Strata without starting the model (#475): the newest code (git pull, when this folder is a git clone), then
|
||||
# what ./setup.sh does before a start - the engine (a new ready-made one when this version needs it, or compiled
|
||||
# again when its source changed), the Python packages, each installed model's settings and draft subset. The model
|
||||
# files are not touched. Start the model later with ./setup.sh. Options are passed on to setup.py.
|
||||
cd "$(dirname "$0")" || exit 1
|
||||
# all of it in a function, read before it runs: the git pull below can change this very file
|
||||
main() {
|
||||
if [ -e .git ]; then
|
||||
if ! command -v git >/dev/null 2>&1; then
|
||||
echo "This folder is a git clone, but git is not installed: install it (sudo apt install git) or run"
|
||||
echo "\"git pull\" here yourself, then run ./update.sh again."
|
||||
exit 1
|
||||
fi
|
||||
echo "Getting the newest Strata (git pull) ..."
|
||||
if ! git pull --ff-only; then
|
||||
echo
|
||||
echo "git pull did not succeed (the reason is above): nothing was updated. Files you changed here can stop it:"
|
||||
echo "\"git status\" lists them."
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
echo "This copy of Strata was not made with git, so it cannot fetch new files itself. Download the newest one:"
|
||||
echo " https://github.com/Niko1221/Strata/archive/refs/heads/main.zip"
|
||||
echo "unzip it anywhere and run ./setup.sh (or ./update.sh) in it: it finds the model files in Strata-data and"
|
||||
echo "sets itself up the same way - nothing big is downloaded again."
|
||||
echo "Checking this copy's engine and settings meanwhile ..."
|
||||
fi
|
||||
exec sh ./setup.sh --update "$@"
|
||||
}
|
||||
main "$@"
|
||||
Reference in new issue
Block a user