Compare commits

...
17 Commits
Author SHA1 Message Date
Niko1221 8f279d82a0 KV streaming: the draft layer keeps its KV in VRAM when no pinned RAM is left (engine 0.1.6)
The MTP drafter's ring over a host copy is set up after the expert arena has pinned what it could; when that
pinned allocation fails the drafter now falls back to its fully resident KV instead of stopping the engine
("mtp: state init failed", seen with Swift 1.5 at 262K). setup requires engine 0.1.6.
2026-09-27 08:34:09 +02:00
Niko1221 0616dfb74e KV streaming for the QSA layers (engine 0.1.5)
From 64K context the KV cache lives in pinned RAM and a VRAM block cache holds the part the sparse attention reads
(--kv-resident N; setup turns it on with 32768 when RAM allows). Blocks the indexer selects are made resident on
the device inside the captured graphs (clock eviction, zero-copy reads of misses); writers write through, so the
attention reads exactly the same values. The freed VRAM goes to the expert cache: Q2_0 at 262K decodes 62.6 tok/s
instead of 50.9. The MTP drafter keeps its window in a ring. New ctest kv_stream_parity. Also wires the opt-in
suffix drafter (--suffix-draft) into the engine.
2026-09-27 08:29:03 +02:00
Niko1221 55e2fb8576 setup: require engine v0.1.4 (IQ3_S) 2026-09-26 23:57:44 +02:00
Niko1221 59b0696208 IQ3_S support: the GSQ-RCO IQ3_S model (engine 0.1.4)
- GPU experts: IQ4_XS gate/up (the IQ3_S file keeps layer 47's routed experts in it), transcribed from llama.cpp's
  vec_dot_iq4_xs_q8_1; checked on all 48 layers against ggml's float reference (native_expert_parity: 0 failures).
- MTP draft head: gather_rows takes rows of any size (the IQ3_S head is Q6_K, 2100-byte rows, not a multiple of 16).
- setup: IQ3_S for the original model (Swift 1.5 has none), 64 GB RAM, context capped at 128K like IQ3_XXS.
- README / DETAILS: IQ3_S speed (RTX 5070, Ryzen 5 7600, 64 GB): 54 tok/s at 4K, 42 at 128K, 374-397 prompt.
2026-09-26 23:55:41 +02:00
Niko1221 0223d32663 setup: update an old ready-made engine when an existing install starts
START-HERE.bat on an installed setup went straight to starting the model, so the MIN_ENGINE check (and the
engine download) only ran when setup was run again with its options.  Now every start checks the installed
engine and replaces it when it is older than MIN_ENGINE.  If it cannot (no internet, the model still running,
no ready-made engine for this GPU) the installed engine is kept and starts as before.
2026-09-26 22:24:29 +02:00
Niko1221 d0017a7c58 setup: require engine v0.1.3 (conversation cache, idle CPU fix, Linux pool fix) 2026-09-26 22:18:20 +02:00
Niko1221 0c2e63e930 Merge pull request #10: Conversation cache + short prompt parts through the decode windows (#8, #10) 2026-09-26 21:25:03 +02:00
Niko1221 888751a0e1 Merge pull request #7: serve: fix zero-token replies from the engine-queue race between requests (#7) 2026-09-26 21:25:03 +02:00
Niko1221 af1c91f245 Merge pull request #9: CPU pool: sleep between requests instead of spinning (#9, fixes #4) 2026-09-26 21:25:03 +02:00
Niko1221 14b22c51c1 Merge pull request #14: CPU pool: fix dangling else in Linux physical_cores() (#14, fixes #13 and #11) 2026-09-26 21:25:03 +02:00
Vistawizard be028948f4 CPU pool: fix dangling else in the Linux physical_cores() branch
Without braces the `else` bound to `if (CPU_ISSET(i, &set))`, so every CPU id
outside the affinity mask appended all hardware_concurrency() ids again:
16143 pool workers on a 16-CPU box, load > 10000 and a hang at start.
Brace both branches so the fallback only runs when sched_getaffinity fails.

Fixes #13, root cause of #11.
2026-09-26 19:23:15 +02:00
Mirtraxxx 470b3d7c21 serve: read short prompt parts through the decode windows
The batched prompt path costs ~300 ms per run however few tokens it
has (it streams every non-resident expert the chunk routes to over
PCIe), plus ~180 ms to refill the expert-cache slots it borrows. So a
one-word message waited ~1.4 s for its first token, and every chat
follow-up paid the same fixed cost.

Each part of the prompt (up to the checkpoint, then the assistant
header) now picks its own path: parts of at most --short-read N tokens
(default 64, 0 = off) go through the verify windows, S tokens at a
time, as decode reads them (~16 ms a token). Longer parts and parts
with picture rows stay batched. Slots are lent just before a batched
run and refilled before a window read, so windows always see the whole
expert cache. STRATA_CKPT_REREAD keeps every read batched.
2026-09-26 01:02:20 -03:00
Mirtraxxx be0e0c94e3 cpu pool: sleep between requests instead of spinning (fixes #4)
A parked expert-pool worker spun on `epoch_` with `_mm_pause` forever, so every
worker's core stayed at 100% while the engine waited for a request: 7 cores
(44% of a 16-thread Ryzen 7 5700X) with nothing to do.

A parked worker now spins as before for up to 20 ms, then blocks on a
condition variable. Every publish goes through `publish()`, which bumps the
epoch and notifies only when a worker is actually asleep (one uncontended
load otherwise). Inside a request the gaps between batches are microseconds,
so the token path never sleeps.
2026-09-25 21:27:53 -03:00
Mirtraxxx 42810ee8bb serve: conversation cache - keep the chat between requests, read only what is new
Every request used to start from session_zero and read its whole prompt again,
so a long chat or an agent loop waited minutes per turn. The engine now keeps
the live session and up to 6 checkpoints of the running state in RAM and reads
only the part of a prompt after the longest one it starts with.

- A checkpoint is the state that cannot be rewound - the 36 GDN recurrences and
  conv histories, the 12 QSA indexer tails, the PLE history (~118 MB). The KV
  cache and the pooled indexer keys are positional (a query reads only cells
  below it, and `dead` for its own block), so they stay in place.
- Taken at the last <|im_start|> of each prompt (the prompt is read in two parts
  there: the history, then the new turn's header) and every 16K fresh tokens.
- Used only when the prompt starts with exactly its tokens and pictures (images
  are compared by a hash of their embeddings and grid); checkpoints the next
  request would overwrite are dropped.
- New lines on stdout: RESUME <n> before reading, PP <pos> <n> <ms> <tok/s> per
  chunk, REUSED <n> after; DONE gains <drafts accepted> <drafts offered>
  <reused>. The server ignores what it does not know; PP lines now count as its
  heartbeat, so keep-alives keep flowing during a long prompt.
- --prompt-cache N (default 6, 0 = off), --prompt-cache-every N, --turn-token ID.
- Checks: STRATA_CKPT_REREAD=1 reads a checkpoint's tokens again instead of
  restoring it; STRATA_STATE_HASH=1 prints a fingerprint of every state part
  after each request.
2026-09-25 21:18:23 -03:00
btechnet 0f8b5d4938 serve: fix zero-token replies from engine-queue race between requests
Symptom: with two requests in flight (opencode fires a title-generation
request alongside the main tool-calling turn), the second request often
returned zero tokens. The server logged `done: 0 tokens in 0 s (length,
cancel=False)` and opencode showed an empty reply.

Cause: StrataEngine multiplexes every request over one resident
`strata --serve` subprocess and a single shared line queue (self.lines),
with no per-request framing. In Service.run(), a stop token (<|im_end|>)
breaks the token loop with the engine generator still suspended and its
internal `done` flag False. That generator's cleanup - write STOP, then
drain self.lines until DONE - therefore ran only when run()'s frame was
torn down later, AFTER the `with self.fifo` lock had been released. By
then the next request held the lock and had issued its own GEN; the prior
request's drain raced it on the shared queue and consumed the NEW
request's DONE. The new request saw an immediate end-of-stream: 0 tokens,
with `finish` left at its default of "length".

Fix: close the engine generator inside the `with self.fifo` block via
try/finally, so a request's STOP+drain completes while it still holds the
lock and cannot leave stray lines on the shared queue for the next
request to misread as its own DONE.

Also in this change:
- Clamp a non-positive max_new to 1024 in both /v1/chat/completions and
  /v1/messages. Some clients send max_tokens: -1 ("unlimited"); that value
  passed straight through and `GEN -1` made the engine yield 0 tokens.
- Distinguish a client disconnect: catch GeneratorExit and set
  finish="disconnect" instead of leaving it as "length".
- Diagnostics gated on STRATA_DEBUG (off by default): a per-request
  `_debug_req` line (api, msg/tool counts, raw max_tokens, resolved
  max_new, thinking flag, prompt tokens) and a `raw:` dump of the decoded
  model output at completion.
- Progress reporting: add tok/s and cancel state to the `done:` line, and
  emit the progress line every 1 s instead of every 15 s for live
  throughput.
2026-09-25 18:11:01 -05:00
Niko1221 07b4a6263a Engine v0.1.2: streaming tool calls, live status, real cancel; fix requests stalling on a full GPU
Streaming
- tool calls stream as they are written (OpenAI tool_calls deltas, Anthropic input_json_delta)
- keep-alives while a long prompt is read, progress in the server window, GET /status
- closing the connection or pressing stop really stops the engine (STOP)

Stalled requests (mostly image turns)
- the expert cache is sized after its slots are written (Windows only counts GPU memory once it is
  touched), so the reserve really stays free instead of the card running at 0 MiB and paging
- the prompt chunk shrinks when its buffers do not fit in the cache, instead of allocating on top
- the image encoder warms up before the engine starts; images are encoded in request order
- the log shows the VRAM left once everything is loaded

setup.py now requires engine 0.1.2 (downloaded automatically on the next START-HERE.bat).
2026-09-25 06:13:46 +02:00
Niko1221 e550970773 README: a short, beginner-friendly front page (details in docs/DETAILS.md); pictures in WebP/TIFF/AVIF 2026-09-24 23:17:29 +02:00
29 changed files with 2774 additions and 492 deletions

No files matched your search

+8 -1
View File
@@ -177,6 +177,7 @@ if(STRATA_ENABLE_CUDA)
src/kernels/cuda/native_qsa_score.cu
src/kernels/cuda/ple.cu
src/kernels/cuda/kv_q8.cu
src/kernels/cuda/kv_stream.cu
src/kernels/cuda/fused_gr.cu
src/kernels/cuda/qsa_decode_attn.cu
src/kernels/cuda/fused_gdn.cu
@@ -250,7 +251,7 @@ if(STRATA_ENABLE_CUDA)
add_executable(strata src/program/generate.cpp)
target_include_directories(strata PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(strata PRIVATE strata_engine strata_prefill CUDA::cudart)
target_link_libraries(strata PRIVATE strata_engine strata_prefill strata_spec CUDA::cudart)
add_executable(dequant_s2_parity src/kernels/dequant_s2_parity.cpp)
# The tests are HOST C++ files that call cudaMalloc/cudaMemcpy themselves, so they need the toolkit headers
@@ -356,6 +357,12 @@ if(STRATA_ENABLE_CUDA)
target_link_libraries(kv_q8_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME kv_q8_parity COMMAND kv_q8_parity)
# KV streaming (docs/kv-streaming-design.md): streamed vs resident attention, bitwise (GPU, synthetic, no model).
add_executable(kv_stream_parity src/kernels/kv_stream_parity.cpp)
target_include_directories(kv_stream_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_link_libraries(kv_stream_parity PRIVATE strata_kernels CUDA::cudart)
add_test(NAME kv_stream_parity COMMAND kv_stream_parity)
# Plan v0.3 P3: native MMVQ for 1..8 tokens (bitwise selftest + --bench timing); synthetic data, no model.
if(STRATA_BUILD_TESTS)
add_executable(native_mmvq_multi bench/micro/native_mmvq_multi.cpp)
+103 -293
View File
@@ -1,332 +1,142 @@
<h1 align="center">Qwen3.8-Flash-Next (125B MoE) on a 12-24 GB NVIDIA GPU + 64 GB RAM: one-click install for Windows / Linux</h1>
<h1 align="center">Strata</h1>
<p align="center"><b>Strata</b>: an inference engine built for exactly one model and exactly this kind of PC</p>
<p align="center"><b>Run a 125-billion-parameter AI model on a normal gaming PC</b><br>
one NVIDIA card (12-24 GB) + 64 GB of RAM · Windows or Linux · one click to install</p>
A 125-billion-parameter mixture-of-experts model (6B active per token) on a **normal gaming PC**: one NVIDIA card with
12-24 GB of VRAM, 64 GB of RAM and a desktop CPU. Strata keeps the always-used part of the model on the GPU, all 24,576
experts in RAM, computes the experts the GPU does not hold on the CPU at the same time, and checks several tokens per
pass with the model's own multi-token-prediction (MTP) layer. The text it produces is identical to plain greedy decoding.
Strata runs **[Qwen3.8-Flash-Next](https://huggingface.co/Qwen/Qwen3.8-Flash-Next)** - a large, smart AI model that
normally needs a server - on your own PC. It writes its answers at **60-95 tokens per second** (a token is about ¾
of a word): faster than you can read.
It serves an **OpenAI-** and **Anthropic-compatible API on localhost**, so any chat app, coding agent or script that
talks to those APIs can use it. Optionally it also **reads images** (screenshots, photos, scanned pages) with the
model's own vision encoder.
- **Private:** everything runs on your PC. Nothing is sent anywhere.
- **Works with your apps:** chat apps, coding agents and scripts that speak the OpenAI or Anthropic API just work.
- **Sees pictures** too, if you want (screenshots, photos, scanned pages).
- **Free and open source.**
> **Jump to:** [Speed](#speed-measured) · [Other GPUs](#other-gpus-estimated) · [Which model?](#which-model) ·
> [Windows](#windows) · [Linux](#linux) · [Using it](#using-it) · [Images](#images-vision) ·
> [Troubleshooting](#troubleshooting) · [How it works](#how-it-works)
> **Jump to:** [Is my PC enough?](#is-my-pc-enough) · [Install](#install-3-steps) · [How fast?](#how-fast-is-it) ·
> [Which model?](#which-model-should-i-pick) · [Using it](#using-it) · [Problems?](#something-went-wrong) ·
> [How it works](#how-does-it-work) · [All the details](docs/DETAILS.md)
---
## Speed (measured)
## Is my PC enough?
RTX 5070 **12 GB**, Ryzen 5 7600 (6 cores), 64 GB DDR5-5200, Windows. One code-agent prompt per length, 256 generated
tokens, MTP speculative decoding on. "262K" is the model's full context window (a 259,943-token prompt).
### Prompt processing (tokens/s)
| Model | 1K | 4K | 32K | 64K | 128K | 262K |
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
| **Q2_0** | 389 | 539 | 571 | 561 | 543 | 496 |
| **IQ2_XS** | 332 | 463 | 495 | 486 | 472 | 437 |
| **IQ3_XXS** | 285 | 410 | 435 | 427 | 414 | - |
### Output (tokens/s)
| Model | 1K | 4K | 32K | 64K | 128K | 262K |
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
| **Q2_0** | 88.7 | 94.6 | 87.5 | 76.4 | 65.1 | 56.3 |
| **IQ2_XS** | 82.0 | 78.0 | 65.3 | 63.7 | 52.0 | 48.0 |
| **IQ3_XXS** | 64.6 | 65.6 | 57.3 | 54.4 | 44.8 | - |
IQ3_XXS at 262K is not measured: with its 43 GB of experts, a 260K-token context brings a 64 GB PC to its memory
limit. Use up to 128K with IQ3_XXS on 64 GB.
Time to first token is prompt length / prompt speed: about 7 s at 4K, 55 s at 32K, 4 minutes at 128K and 9 minutes at
262K. The raw numbers: `bench/results/`. The [paper](docs/paper/Strata-Paper.pdf) explains every number.
## Other GPUs (estimated)
Not measured - estimated from the runs above (same CPU and 64 GB RAM): the GPU part scaled by memory bandwidth, the CPU
part by how many more experts the card's VRAM holds. Treat as **±20%**. Numbers are *prompt / output* tokens/s.
| GPU | Model | 1K | 4K | 32K | 64K | 128K | 262K |
| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| RTX 5060 Ti 16GB | Q2_0 | ~341 / ~80 | ~472 / ~87 | ~501 / ~81 | ~492 / ~72 | ~476 / ~62 | ~435 / ~53 |
| | IQ2_XS | ~291 / ~80 | ~406 / ~77 | ~434 / ~63 | ~426 / ~62 | ~413 / ~51 | ~383 / ~47 |
| | IQ3_XXS | ~249 / ~66 | ~359 / ~65 | ~381 / ~56 | ~374 / ~54 | ~363 / ~45 | - |
| RTX 3090 24GB | Q2_0 | ~355 / ~128 | ~491 / ~140 | ~521 / ~130 | ~512 / ~115 | ~495 / ~100 | ~453 / ~85 |
| | IQ2_XS | ~303 / ~131 | ~422 / ~128 | ~451 / ~103 | ~444 / ~102 | ~430 / ~85 | ~398 / ~78 |
| | IQ3_XXS | ~260 / ~106 | ~374 / ~103 | ~396 / ~89 | ~390 / ~85 | ~378 / ~71 | - |
More VRAM matters more than a faster GPU: every extra GB holds ~700 more experts, and every expert on the GPU is one the
CPU does not have to compute. A 3090's 24 GB takes most of the CPU work away.
## Which model?
All three are [ISTA-DASLab's GSQ-RCO quantizations](https://huggingface.co/ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF)
of [Qwen3.8-Flash-Next](https://huggingface.co/Qwen/Qwen3.8-Flash-Next).
| Model | Download | RAM it uses | Speed | Quality |
| --- | ---: | ---: | --- | --- |
| **Q2_0** | 66 GB | ~34 GB experts + ~6 GB | fastest | good |
| **IQ2_XS** | 68 GB | ~36 GB experts + ~6 GB | close to Q2_0 | a bit better |
| **IQ3_XXS** | 76 GB | ~43 GB experts + ~6 GB | slower (more CPU work) | best |
With 64 GB of RAM all three fit (close the browser for IQ3_XXS, and keep its context at 128K or less). With 48 GB only Q2_0 / IQ2_XS may fit. 32 GB is not enough.
### Or: Swift 1.5 (a fine-tune that thinks shorter)
The setup's first question also offers **[Swift 1.5](https://huggingface.co/ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF)**,
UkisAI's fine-tune of Qwen3.8-Flash-Next, trained to reach the answer with much less thinking (its authors: 63% fewer
thinking tokens, 1.8x sooner answers, under 1% accuracy loss). Same architecture, the same three sizes, its own
vision encoder; Strata runs it at the same speed (4K, IQ2_XS: 465 prompt / 78.7 output tokens/s, vs 467 / 78.3 for
the original). Its authors recommend **IQ2_XS** (their Q2_0 is marked experimental). Its license is the Swift Open
License 1.0 - read it on the model page.
Our small check (8 reasoning questions, default thinking, IQ2_XS): both models got **8/8**; Swift used **1,234**
output tokens in 28 s, the original **2,682** in 46 s - most of the difference from one question the original
thought about for 1,524 tokens. Not a benchmark, but consistent with the claim.
```
START-HERE.bat --setup --family swift --model IQ2_XS
```
## Before you start
You need **only an NVIDIA driver** (version 580 or newer; update it with the NVIDIA App or from
[nvidia.com/drivers](https://www.nvidia.com/drivers)). Everything else is installed for you the first time.
| | |
| You need | |
| --- | --- |
| GPU | NVIDIA **RTX 30, 40 or 50 series**, **12 GB VRAM or more** (8 GB runs, slowly). Measured on an RTX 5070; RTX 30/40 are untested. |
| RAM | **64 GB** recommended (see the table above). |
| CPU | x86-64 with AVX2 (any Intel/AMD desktop CPU from the last ~8 years). AVX-512 (Ryzen 7000/9000) is a bit faster. |
| Disk | ~70-80 GB free for the model, ~6 GB for the MTP layer (+1 GB with images). **Q2_0 on an AVX-512 CPU** also writes a one-time ~40 GB copy of its experts for the fast CPU kernel. An NVMe SSD is strongly recommended. |
| OS | Windows 10/11, or Linux (Ubuntu 22.04/24.04 get everything installed automatically). |
| **Graphics card** | NVIDIA RTX 30, 40 or 50 series with **12 GB of VRAM or more** |
| **Memory (RAM)** | **64 GB** |
| **Free disk space** | ~80 GB (an SSD makes the first start much faster) |
| **System** | Windows 10/11, or Linux |
What the first start installs, all inside this folder (`.venv/`, `engine/`, `third_party/`, `models/`, `packs/`, `mtp/`):
Python 3.12 if you have none (for your user account, no admin), a private Python environment, NVIDIA's CUDA libraries
(from pip, ~0.4 GB), the ready-made Strata engine for RTX 30/40/50, the model and the MTP draft layer. If no
ready-made engine fits your PC, it offers to install the build tools (Visual Studio Build Tools + CUDA Toolkit on
Windows, `build-essential` + CUDA on Ubuntu) and compiles the engine for your GPU (asks first; 20-40 minutes once).
That's it. The only thing you install yourself is a current **NVIDIA driver**
([nvidia.com/drivers](https://www.nvidia.com/drivers) or the NVIDIA App). Everything else - Python, the engine, the
model - is set up for you.
---
## Install (3 steps)
## Windows
**Windows**
### Double-click `START-HERE.bat`
1. [Download this project](https://github.com/Niko1221/Strata/archive/refs/heads/main.zip) and unzip it (or `git clone` it).
2. Double-click **`START-HERE.bat`**.
3. Answer 4 questions - or just press Enter each time for the recommended choice:
- **Which model?** The original, or Swift 1.5 (a version that thinks shorter and answers sooner)
- **Which size?** Q2_0, IQ2_XS, IQ3_XXS or IQ3_S - see [which model](#which-model-should-i-pick)
- **How much context?** How much text it can keep in mind at once (it suggests one for your card)
- **Images?** Whether it should also read pictures
**The first time** it asks four questions and does the rest:
Then it downloads everything (the model is ~70 GB, so the first time takes a while - you can stop and it picks up
where it left off) and **starts the model**. Your browser opens a chat page at `http://127.0.0.1:8080`.
1. **Which model?** Qwen3.8-Flash-Next (the original) or Swift 1.5 (the fine-tune that thinks shorter).
2. **Which size?** Q2_0, IQ2_XS or IQ3_XXS (it recommends one for your RAM).
3. **How much context?** 8K to 256K tokens (it recommends one for your VRAM).
4. **Images?** yes / no (see [Images](#images-vision)).
**Next time**, just double-click `START-HERE.bat` again: it starts right away, nothing is downloaded twice. Close its
window to stop the model.
Then it downloads and prepares everything (the model is 66-76 GB, so the first start takes a while; an interrupted
download continues where it stopped) and **starts the model**: your browser opens `http://127.0.0.1:8080`, a small
page that shows it is running and lets you chat. The API is at `http://127.0.0.1:8080/v1` for your apps.
**Linux:** run `./setup.sh` - same questions, same result.
**Every time after that**, `START-HERE.bat` just starts the model (30-90 s to load 34-43 GB into RAM). Nothing is
downloaded again. Closing the window stops the model.
## How fast is it?
```
START-HERE.bat --setup install another model, or change context / images
START-HERE.bat --model IQ2_XS --context 32768 --vision yes --yes no questions
START-HERE.bat --gguf-dir D:\models\IQ2_XS use GGUF files you already have
START-HERE.bat --port 8081 another port
```
Measured on an RTX 5070 (12 GB), a Ryzen 5 7600 and 64 GB of RAM:
With more than one model installed, it asks which one to start. `run-<model>.bat` starts a model directly.
| Size | Writes answers (short chat) | Writes answers (128K context) | Reads your prompt |
| --- | ---: | ---: | ---: |
| **Q2_0** | 95 tokens/s | 65 tokens/s | 539 tokens/s |
| **IQ2_XS** | 78 tokens/s | 52 tokens/s | 463 tokens/s |
| **IQ3_XXS** | 66 tokens/s | 45 tokens/s | 410 tokens/s |
| **IQ3_S** | 54 tokens/s | 42 tokens/s | 374 tokens/s |
### Chat in the terminal (optional)
- **Writes answers** = how fast the reply appears (tokens per second).
- **Reads your prompt** = how fast it takes in what you send (long documents, code, chat history).
```
.venv\Scripts\python chat.py
```
A card with more VRAM is faster, because more of the model fits on the GPU: an RTX 3090 (24 GB) should do roughly
100-140 tokens per second. All measurements, long-context numbers and estimates for other cards are in the
[details](docs/DETAILS.md#speed-measured).
---
## Which model should I pick?
## Linux
**The size** (the same model, compressed more or less):
```bash
./setup.sh
```
| Size | Download | Speed | Quality | Pick it if... |
| --- | ---: | --- | --- | --- |
| **Q2_0** | 66 GB | fastest | good | you want speed |
| **IQ2_XS** | 68 GB | fast | better | you want a good all-rounder (**recommended**) |
| **IQ3_XXS** | 76 GB | slower | great | you want better answers (uses 43 GB of your 64 GB RAM) |
| **IQ3_S** | 84 GB | slowest | best: matches the full model on the published tests | you want the very best answers (original model only; uses 50 GB of your 64 GB RAM, so close other big programs) |
The same questions, the same automatic install (it uses `sudo apt` for Python and, only if it has to compile,
for the build tools), and the same start: `http://127.0.0.1:8080`. Later runs of `./setup.sh` (or `./run-<model>.sh`)
start the model directly. Options as on Windows (`./setup.sh --setup`, `--model Q2_0 --yes`, `--gguf-dir /data/Q2_0`).
Terminal chat: `.venv/bin/python chat.py`.
**The version:**
---
- **Qwen3.8-Flash-Next** - the original.
- **[Swift 1.5](https://huggingface.co/ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF)** - a fine-tune by UkisAI
that thinks much shorter before answering, so you get the answer sooner, with about the same quality. Same speed per
token. Its own license applies (see its page).
Not sure? Take **IQ2_XS**. You can add another one later with `START-HERE.bat --setup`.
## Using it
The server listens on `http://127.0.0.1:8080` (change with `--port` in setup, or edit the run script).
- **Chat in the browser:** `http://127.0.0.1:8080` - a simple chat page (it opens by itself when the model starts).
- **Chat in the terminal:** `.venv\Scripts\python chat.py`
- **Your apps and coding agents:** add it as an "OpenAI-compatible" provider with base URL
**`http://127.0.0.1:8080/v1`**, any API key and any model name. Apps that use Anthropic's API: `http://127.0.0.1:8080/v1/messages`.
- **Thinking:** the model thinks before it answers. Choose **off, low, medium or high** - in the chat page menu, with
`/think low` in `chat.py`, or with your app's "reasoning effort" setting. Off is fastest; high is best for hard questions.
- **Pictures:** in the chat page click **Picture**; in `chat.py` type `/image <path>`; in apps just attach them.
- **From your phone or another PC:** see the [details](docs/DETAILS.md#using-it) (set an API key first).
| API | Endpoint |
**Good to know:** it answers one request at a time. The first message of a chat is read in full (about 1 minute per
30,000 tokens); after that it keeps the conversation and reads only what is new, so follow-ups start in seconds.
## Something went wrong?
| What you see | What to do |
| --- | --- |
| OpenAI Chat Completions (stream and non-stream, tools) | `POST /v1/chat/completions` |
| Anthropic Messages (stream and non-stream, tools) | `POST /v1/messages` |
| Model list / health | `GET /v1/models`, `GET /health` |
| `the NVIDIA driver is too old` | Update the driver (NVIDIA App or nvidia.com/drivers), restart the PC, run `START-HERE.bat` again. |
| It stopped during download or setup | Run `START-HERE.bat` again - it continues where it stopped. |
| `port 8080 is already in use` | Strata is already running - look for its window. |
| The first start takes minutes | Normal: it loads 35-43 GB into RAM. The next start is faster. |
| Slow, and the disk light is busy | Not enough free RAM: close other programs (browsers use a lot), or pick Q2_0 / IQ2_XS. |
| "prompt exceeds the context" | The conversation is longer than the context you chose: run `START-HERE.bat --setup` and pick more. |
```bash
curl http://127.0.0.1:8080/v1/chat/completions -H "Content-Type: application/json" -d '{
"model": "strata", "messages": [{"role": "user", "content": "Write a haiku about GPUs."}], "max_tokens": 512 }'
```
More in the [full troubleshooting table](docs/DETAILS.md#troubleshooting). Still stuck? Open an issue and attach
`strata-<model>.log` from this folder.
```python
from openai import OpenAI
client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="none")
r = client.chat.completions.create(model="strata", messages=[{"role": "user", "content": "Hello!"}])
print(r.choices[0].message.content)
```
## How does it work?
- **Thinking levels: none, low, medium, high.** The model thinks before it answers (streamed as
`reasoning_content`, Anthropic: `thinking` blocks). Choose how much per request - in the chat page (the "Thinking"
menu), in `chat.py` (`/think low`), or over the API:
A model this big doesn't fit on a gaming graphics card. Strata splits the work between the parts of your PC:
| API | how |
| --- | --- |
| OpenAI | `"reasoning_effort": "none" \| "low" \| "medium" \| "high"` (also `"reasoning": {"effort": ...}`, or `"chat_template_kwargs": {"enable_thinking": false}`) |
| Anthropic | `"output_config": {"effort": "low" \| "medium" \| "high"}`, `"thinking": {"type": "disabled"}`, or `"thinking": {"type": "enabled", "budget_tokens": N}` (under 2K = low, under 8K = medium, more = high) |
<p align="center"><img src="docs/paper/tiers.svg" width="700" alt="how Strata splits the model between GPU, RAM and SSD"></p>
Without a setting the model uses its own default, **high**. `none` answers at once (fastest); `low` keeps the thinking
short. The levels are instructions the model was trained with, not a hard token limit: on easy questions all three
think briefly, on hard ones `high` thinks longest and is most accurate.
- **Chat apps.** Any app with an "OpenAI-compatible" provider works: base URL `http://127.0.0.1:8080/v1`, any API key.
- **Context.** Chosen in setup (8K-262K). Requests longer than that are refused, never silently cut.
- **From other devices / the internet.** The server listens on your PC only (`127.0.0.1`). To reach it from elsewhere,
put a tunnel in front of it, for example [cloudflared](https://developers.cloudflare.com/cloudflare-one/connections/connect-networks/do-more-with-tunnels/trycloudflare/):
`cloudflared tunnel --url http://127.0.0.1:8080`. **Set a key first**, or anyone with the link can use your PC:
add `"api_key": "some-long-secret"` to `strata-<model>.json` (or set the `STRATA_API_KEY` environment variable);
clients then send it as their API key.
- **The GPU** runs the part of the model that is used for every word, plus the "experts" it needs most often.
- **The RAM** holds all 24,576 experts, and **the CPU** computes the few the GPU doesn't have - at the same time as the GPU.
- **The SSD** holds a big lookup table; the model reads a few rows of it per word.
- **A small helper inside the model guesses the next words**, and Strata checks several guesses at once. That makes
it 1.6-1.8x faster than going word by word - and the answer is exactly the same.
**Current limits (v1):** one request at a time; greedy decoding (temperature is ignored); every request processes its
whole prompt again (no conversation cache yet, so long chats have a long time-to-first-token); images only when set up
with them (below); no video.
The full story is in the [paper](docs/paper/Strata-Paper.pdf) and the [details](docs/DETAILS.md).
---
## Credits
## Images (vision)
The model has a vision encoder: [`mmproj-Qwen3.8-Flash-Next-BF16.gguf`](https://huggingface.co/ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF)
(0.9 GB, a 27-layer ViT plus the projector into the language model). It is **optional**: say yes when the setup asks
"Images?", or run it again with `--vision gpu` (or `--vision cpu`). The setup downloads the encoder, builds a small
helper (`strata-vision`, from llama.cpp's `mtmd` library) and adds it to your start script. Nothing else changes.
| Encoder on | Time per picture | Cost |
| --- | --- | --- |
| **GPU** (recommended) | **0.1-0.5 s** (up to 1,024 image tokens) | ~1.4 GB of VRAM is kept free for it, so the expert cache is smaller: text output is a few % slower (table below) |
| CPU | 10-30 s (pictures are scaled down to ~300 image tokens) | nothing on the GPU |
A picture becomes up to 1,024 tokens of the context (a 640x480 photo: 300). The same picture sent again, as chat apps
do on every turn, is encoded only once.
### Sending a picture
**Terminal chat:** type `/image <path to a picture>`, press Enter, then type your question.
```
you> /image C:\Users\me\Pictures\receipt.jpg
(picture attached: receipt.jpg - now type your question)
you> What is the total on this receipt?
```
**OpenAI API** (an `image_url` part: a `data:` URL, an `http(s)://` URL or a local file path):
```python
import base64
from openai import OpenAI
client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="none")
img = base64.b64encode(open("photo.jpg", "rb").read()).decode()
r = client.chat.completions.create(model="strata", messages=[{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img}"}},
{"type": "text", "text": "What is in this picture?"}]}])
print(r.choices[0].message.content)
```
**Anthropic API:** an `image` block with a `base64` (or `url`) source, as usual.
JPEG, PNG, BMP and GIF work (WebP does not). Chat apps with image upload work the same way.
### Speed with images on (4K context, measured)
| Model | Prompt tok/s, images off | Prompt tok/s, images on | Output tok/s, images off | Output tok/s, images on |
| --- | ---: | ---: | ---: | ---: |
| **Q2_0** | 541 | 531 | 90.7 | 86.9 |
| **IQ2_XS** | 468 | 458 | 77.1 | 74.0 |
| **IQ3_XXS** | 411 | 401 | 64.7 | 59.5 |
"Off" is the published setup; "on" runs with the encoder loaded on the GPU and its VRAM kept free, so ~1,000 fewer
experts fit in VRAM and the CPU computes a few more per token: 2-8% slower. For text, turning images on changes
nothing else (the output is bit-identical when the VRAM is the same).
A question about a picture (a 640x480 newspaper page = 300 image tokens, a 328-token prompt; the whole request, encoder
on the GPU, measured through the API):
| Model | Answer ("MEN WALK ON MOON") | Picture encoded | Prompt | Output |
| --- | ---: | ---: | ---: | ---: |
| **Q2_0** | 1.9 s | ~0.1 s | 207 tok/s | 65-73 tok/s |
| **IQ2_XS** | 2.3 s | ~0.1 s | 173 tok/s | 50-60 tok/s |
| **IQ3_XXS** | 2.7 s | ~0.1 s | 144 tok/s | 41-50 tok/s |
(Short prompts run below the 4K prompt speed: a 2,048-token chunk is where the prompt path is efficient. Short answers
run below the long-output speed: the first rounds have no draft yet.)
**How it works inside:** the encoder turns the picture into rows of the same width as the model's word embeddings;
Strata puts them where the prompt has `<|image_pad|>` tokens and gives each one its 2-D position (row and column in the
picture; the model uses interleaved M-RoPE). Answers match llama.cpp's multimodal implementation token for token on our
test images.
---
## Troubleshooting
| Symptom | What to do |
| --- | --- |
| `the NVIDIA driver is too old` | Update the driver (NVIDIA App or nvidia.com/drivers), restart, run `START-HERE.bat` again. |
| Python or the build tools could not be installed | Install what it names (links are printed), then run it again. Everything already done is kept. |
| `port 8080 is already in use` | Strata is already running (look for its window), or another program uses the port: `START-HERE.bat --port 8081`. |
| `cudaHostRegister ... out of memory` in the log | Normal on Windows: the engine pins the experts in per-layer slices instead. Only a problem if the load then fails. |
| The first start takes minutes | It is reading 34-43 GB into RAM; the second start is faster while the files are in the OS cache. |
| Slow output, disk light busy | Not enough free RAM: close other programs, or choose Q2_0 / IQ2_XS. |
| `prompt ... exceeds the context` | The request is longer than the context you chose: run setup again with a bigger `--context`. |
| Slower than the tables | The monitor plugged into the GPU and other GPU programs take VRAM from the expert cache; RAM running below its rated speed (enable EXPO/XMP in the BIOS) slows the CPU half. |
| `this server was started without the vision encoder` | The model was set up for text only: run setup again with `--vision gpu`. |
| A picture is refused or `cannot read the image` | Use JPEG, PNG, BMP or GIF (not WebP). |
| Pictures are slow (10-30 s) | The encoder runs on the CPU: run setup again with `--vision gpu` (needs ~1.4 GB of VRAM). |
| Anything else | The engine log is `strata-<model>.log` in this folder. |
---
## How it works
<p align="center"><img src="docs/paper/tiers.svg" width="760" alt="memory tiers"></p>
- **GPU (VRAM):** attention and DeltaNet mixers, the gated-residual weights, routers, shared experts, output head, the MTP
draft layer, the KV cache, and an **expert cache** that fills the rest of VRAM with the most-used experts (it adapts to
the conversation while you chat).
- **RAM:** all 24,576 experts, pinned. The CPU computes the experts that are not on the GPU **in place**, at the same time
as the GPU works on the cached ones (AVX-512 / AVX2 kernels, ggml's for the i-quants).
- **SSD:** the 28.8 GB n-gram table, read a few rows per token through the OS cache.
- **Speculation:** the model's own MTP layer drafts up to 3 tokens; one pass over all 48 layers checks them. 2.4-3.2
tokens per pass on average.
- **Prompts** are processed in 2,048-token chunks with the experts streamed to the GPU over PCIe.
The full story, with measurements, bottlenecks and what comes next: **[docs/paper/Strata-Paper.pdf](docs/paper/Strata-Paper.pdf)**.
---
## Credits and licenses
- Model: [Qwen/Qwen3.8-Flash-Next](https://huggingface.co/Qwen/Qwen3.8-Flash-Next) by the Qwen team; quantizations:
[ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF](https://huggingface.co/ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF).
Swift 1.5: [ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF](https://huggingface.co/ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF)
by UkisAI. Their licenses apply to the weights.
- [llama.cpp / ggml](https://github.com/ggml-org/llama.cpp) (MIT): the i-quant formats, the GPU dot products and
dequantizers transcribed in `src/kernels/cuda/iq_kernels.cu`, the CPU backend linked for the i-quant experts, the
`mtmd` library behind the image encoder (`tools/vision/`), and `gguf-py` used by the tools. See
`third_party/ggml/LICENSE`.
- Ideas from [Splash](https://github.com/incoai/splash), [ninfer](https://github.com/Neroued/ninfer) and
[HyperQwen](https://github.com/syv-ai/HyperQwen); references in the paper.
- Model: [Qwen3.8-Flash-Next](https://huggingface.co/Qwen/Qwen3.8-Flash-Next) by the Qwen team; compressed versions by
[ISTA-DASLab](https://huggingface.co/ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF);
[Swift 1.5](https://huggingface.co/ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF) by UkisAI. Their licenses apply
to the model files.
- Built with parts of [llama.cpp / ggml](https://github.com/ggml-org/llama.cpp) (MIT). Ideas from
[Splash](https://github.com/incoai/splash), [ninfer](https://github.com/Neroued/ninfer) and
[HyperQwen](https://github.com/syv-ai/HyperQwen). More in the [details](docs/DETAILS.md#credits-and-licenses).
+346
View File
@@ -0,0 +1,346 @@
# Strata - the details
The technical side of Strata: every measured number, the API, images, all settings and how the engine works.
New here? Start with the [README](../README.md) - it has everything you need to install and use it.
> **On this page:** [Speed](#speed-measured) · [Other GPUs](#other-gpus-estimated) · [Which model?](#which-model) ·
> [Requirements](#before-you-start) · [Windows](#windows) · [Linux](#linux) · [API](#using-it) · [Images](#images-vision) ·
> [Troubleshooting](#troubleshooting) · [How it works](#how-it-works)
---
## Speed (measured)
RTX 5070 **12 GB**, Ryzen 5 7600 (6 cores), 64 GB DDR5-5200, Windows. One code-agent prompt per length, 256 generated
tokens, MTP speculative decoding on. "262K" is the model's full context window (a 259,943-token prompt).
### Prompt processing (tokens/s)
| Model | 1K | 4K | 32K | 64K | 128K | 262K |
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
| **Q2_0** | 389 | 539 | 571 | 561 | 543 | 496 |
| **IQ2_XS** | 332 | 463 | 495 | 486 | 472 | 437 |
| **IQ3_XXS** | 285 | 410 | 435 | 427 | 414 | - |
| **IQ3_S** | 260 | 374 | 397 | - | 378 | - |
### Output (tokens/s)
| Model | 1K | 4K | 32K | 64K | 128K | 262K |
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
| **Q2_0** | 88.7 | 94.6 | 87.5 | 76.4 | 65.1 | 56.3 |
| **IQ2_XS** | 82.0 | 78.0 | 65.3 | 63.7 | 52.0 | 48.0 |
| **IQ3_XXS** | 64.6 | 65.6 | 57.3 | 54.4 | 44.8 | - |
| **IQ3_S** | 51.2 | 54.4 | 51.1 | - | 42.2 | - |
IQ3_XXS and IQ3_S at 262K are not measured: with their 43 / 50 GB of experts, a 260K-token context brings a 64 GB PC
to its memory limit. Use up to 128K with them on 64 GB (setup caps it). IQ3_S (engine 0.1.4 or newer) is only published
for the original model, not for Swift 1.5.
**KV streaming (engine 0.1.5):** at 64K and more, setup keeps the context's KV cache in RAM and only the part the
attention reads in VRAM (`--kv-resident 32768`), so more experts fit on the GPU. Q2_0 at 262K: 50.9 -> 62.6 tokens/s
(1,589 -> 3,872 experts in VRAM); at 128K about +6%. The attention reads exactly the same values (only where the KV lives changes); it
costs ~13.7 KB of RAM per context token (1.7 GB at 128K). Existing installs: run `START-HERE.bat --setup` once to turn
it on.
Time to first token is prompt length / prompt speed: about 7 s at 4K, 55 s at 32K, 4 minutes at 128K and 9 minutes at
262K. The raw numbers: [`bench/results/`](../bench/results/). The [paper](paper/Strata-Paper.pdf) explains every number.
## Other GPUs (estimated)
Not measured - estimated from the runs above (same CPU and 64 GB RAM): the GPU part scaled by memory bandwidth, the CPU
part by how many more experts the card's VRAM holds. Treat as **±20%**. Numbers are *prompt / output* tokens/s.
| GPU | Model | 1K | 4K | 32K | 64K | 128K | 262K |
| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| RTX 5060 Ti 16GB | Q2_0 | ~341 / ~80 | ~472 / ~87 | ~501 / ~81 | ~492 / ~72 | ~476 / ~62 | ~435 / ~53 |
| | IQ2_XS | ~291 / ~80 | ~406 / ~77 | ~434 / ~63 | ~426 / ~62 | ~413 / ~51 | ~383 / ~47 |
| | IQ3_XXS | ~249 / ~66 | ~359 / ~65 | ~381 / ~56 | ~374 / ~54 | ~363 / ~45 | - |
| RTX 3090 24GB | Q2_0 | ~355 / ~128 | ~491 / ~140 | ~521 / ~130 | ~512 / ~115 | ~495 / ~100 | ~453 / ~85 |
| | IQ2_XS | ~303 / ~131 | ~422 / ~128 | ~451 / ~103 | ~444 / ~102 | ~430 / ~85 | ~398 / ~78 |
| | IQ3_XXS | ~260 / ~106 | ~374 / ~103 | ~396 / ~89 | ~390 / ~85 | ~378 / ~71 | - |
More VRAM matters more than a faster GPU: every extra GB holds ~700 more experts, and every expert on the GPU is one the
CPU does not have to compute. A 3090's 24 GB takes most of the CPU work away.
## Which model?
All three are [ISTA-DASLab's GSQ-RCO quantizations](https://huggingface.co/ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF)
of [Qwen3.8-Flash-Next](https://huggingface.co/Qwen/Qwen3.8-Flash-Next).
| Model | Download | RAM it uses | Speed | Quality |
| --- | ---: | ---: | --- | --- |
| **Q2_0** | 66 GB | ~34 GB experts + ~6 GB | fastest | good |
| **IQ2_XS** | 68 GB | ~36 GB experts + ~6 GB | close to Q2_0 | a bit better |
| **IQ3_XXS** | 76 GB | ~43 GB experts + ~6 GB | slower (more CPU work) | best |
With 64 GB of RAM all three fit (close the browser for IQ3_XXS, and keep its context at 128K or less). With 48 GB only Q2_0 / IQ2_XS may fit. 32 GB is not enough.
### Or: Swift 1.5 (a fine-tune that thinks shorter)
The setup's first question also offers **[Swift 1.5](https://huggingface.co/ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF)**,
UkisAI's fine-tune of Qwen3.8-Flash-Next, trained to reach the answer with much less thinking (its authors: 63% fewer
thinking tokens, 1.8x sooner answers, under 1% accuracy loss). Same architecture, the same three sizes, its own
vision encoder; Strata runs it at the same speed (4K, IQ2_XS: 465 prompt / 78.7 output tokens/s, vs 467 / 78.3 for
the original). Its authors recommend **IQ2_XS** (their Q2_0 is marked experimental). Its license is the Swift Open
License 1.0 - read it on the model page.
Our small check (8 reasoning questions, default thinking, IQ2_XS): both models got **8/8**; Swift used **1,234**
output tokens in 28 s, the original **2,682** in 46 s - most of the difference from one question the original
thought about for 1,524 tokens. Not a benchmark, but consistent with the claim.
```
START-HERE.bat --setup --family swift --model IQ2_XS
```
## Before you start
You need **only an NVIDIA driver** (version 580 or newer; update it with the NVIDIA App or from
[nvidia.com/drivers](https://www.nvidia.com/drivers)). Everything else is installed for you the first time.
| | |
| --- | --- |
| GPU | NVIDIA **RTX 30, 40 or 50 series**, **12 GB VRAM or more** (8 GB runs, slowly). Measured on an RTX 5070; RTX 30/40 are untested. |
| RAM | **64 GB** recommended (see the table above). |
| CPU | x86-64 with AVX2 (any Intel/AMD desktop CPU from the last ~8 years). AVX-512 (Ryzen 7000/9000) is a bit faster. |
| Disk | ~70-80 GB free for the model, ~6 GB for the MTP layer (+1 GB with images). **Q2_0 on an AVX-512 CPU** also writes a one-time ~40 GB copy of its experts for the fast CPU kernel. An NVMe SSD is strongly recommended. |
| OS | Windows 10/11, or Linux (Ubuntu 22.04/24.04 get everything installed automatically). |
What the first start installs, all inside this folder (`.venv/`, `engine/`, `third_party/`, `models/`, `packs/`, `mtp/`):
Python 3.12 if you have none (for your user account, no admin), a private Python environment, NVIDIA's CUDA libraries
(from pip, ~0.4 GB), the ready-made Strata engine for RTX 30/40/50, the model and the MTP draft layer. If no
ready-made engine fits your PC, it offers to install the build tools (Visual Studio Build Tools + CUDA Toolkit on
Windows, `build-essential` + CUDA on Ubuntu) and compiles the engine for your GPU (asks first; 20-40 minutes once).
---
## Windows
### Double-click `START-HERE.bat`
**The first time** it asks four questions and does the rest:
1. **Which model?** Qwen3.8-Flash-Next (the original) or Swift 1.5 (the fine-tune that thinks shorter).
2. **Which size?** Q2_0, IQ2_XS or IQ3_XXS (it recommends one for your RAM).
3. **How much context?** 8K to 256K tokens (it recommends one for your VRAM).
4. **Images?** yes / no (see [Images](#images-vision)).
Then it downloads and prepares everything (the model is 66-76 GB, so the first start takes a while; an interrupted
download continues where it stopped) and **starts the model**: your browser opens `http://127.0.0.1:8080`, a small
page that shows it is running and lets you chat. The API is at `http://127.0.0.1:8080/v1` for your apps.
**Every time after that**, `START-HERE.bat` just starts the model (30-90 s to load 34-43 GB into RAM). Nothing is
downloaded again. Closing the window stops the model.
```
START-HERE.bat --setup install another model, or change context / images
START-HERE.bat --model IQ2_XS --context 32768 --vision yes --yes no questions
START-HERE.bat --gguf-dir D:\models\IQ2_XS use GGUF files you already have
START-HERE.bat --port 8081 another port
```
With more than one model installed, it asks which one to start. `run-<model>.bat` starts a model directly.
### Chat in the terminal (optional)
```
.venv\Scripts\python chat.py
```
---
## Linux
```bash
./setup.sh
```
The same questions, the same automatic install (it uses `sudo apt` for Python and, only if it has to compile,
for the build tools), and the same start: `http://127.0.0.1:8080`. Later runs of `./setup.sh` (or `./run-<model>.sh`)
start the model directly. Options as on Windows (`./setup.sh --setup`, `--model Q2_0 --yes`, `--gguf-dir /data/Q2_0`).
Terminal chat: `.venv/bin/python chat.py`.
---
## Using it
The server listens on `http://127.0.0.1:8080` (change with `--port` in setup, or edit the run script).
| API | Endpoint |
| --- | --- |
| OpenAI Chat Completions (stream and non-stream, tools) | `POST /v1/chat/completions` |
| Anthropic Messages (stream and non-stream, tools) | `POST /v1/messages` |
| Model list / health | `GET /v1/models`, `GET /health` |
| What the model is doing right now | `GET /status` |
```bash
curl http://127.0.0.1:8080/v1/chat/completions -H "Content-Type: application/json" -d '{
"model": "strata", "messages": [{"role": "user", "content": "Write a haiku about GPUs."}], "max_tokens": 512 }'
```
```python
from openai import OpenAI
client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="none")
r = client.chat.completions.create(model="strata", messages=[{"role": "user", "content": "Hello!"}])
print(r.choices[0].message.content)
```
- **Thinking levels: none, low, medium, high.** The model thinks before it answers (streamed as
`reasoning_content`, Anthropic: `thinking` blocks). Choose how much per request - in the chat page (the "Thinking"
menu), in `chat.py` (`/think low`), or over the API:
| API | how |
| --- | --- |
| OpenAI | `"reasoning_effort": "none" \| "low" \| "medium" \| "high"` (also `"reasoning": {"effort": ...}`, or `"chat_template_kwargs": {"enable_thinking": false}`) |
| Anthropic | `"output_config": {"effort": "low" \| "medium" \| "high"}`, `"thinking": {"type": "disabled"}`, or `"thinking": {"type": "enabled", "budget_tokens": N}` (under 2K = low, under 8K = medium, more = high) |
Without a setting the model uses its own default, **high**. `none` answers at once (fastest); `low` keeps the thinking
short. The levels are instructions the model was trained with, not a hard token limit: on easy questions all three
think briefly, on hard ones `high` thinks longest and is most accurate.
- **Streaming.** With `"stream": true` everything arrives as it is made: the thinking, the answer, and tool calls
(the tool's name first, then its arguments piece by piece, like OpenAI and Anthropic do). While the model reads a
long prompt the stream sends keep-alives, so agents do not time out; the server window prints progress every
15 s, and `GET /status` says what it is doing (`reading the prompt`, `answering`, tokens so far). Closing the
connection or pressing stop in your app really stops the model, so the next request starts at once.
- **Chat apps.** Any app with an "OpenAI-compatible" provider works: base URL `http://127.0.0.1:8080/v1`, any API key.
- **Context.** Chosen in setup (8K-262K). Requests longer than that are refused, never silently cut.
- **From other devices / the internet.** The server listens on your PC only (`127.0.0.1`). To reach it from elsewhere,
put a tunnel in front of it, for example [cloudflared](https://developers.cloudflare.com/cloudflare-one/connections/connect-networks/do-more-with-tunnels/trycloudflare/):
`cloudflared tunnel --url http://127.0.0.1:8080`. **Set a key first**, or anyone with the link can use your PC:
add `"api_key": "some-long-secret"` to `strata-<model>.json` (or set the `STRATA_API_KEY` environment variable);
clients then send it as their API key.
**Conversation cache.** A request that continues a chat reads only the part after what the engine already holds: the
live session, or one of the checkpoints it keeps in RAM (up to 6, ~118 MB each, taken at the start of each new
assistant turn and every 16K prompt tokens). A checkpoint is used only when the prompt starts with exactly its tokens
and pictures. Engine options: `--prompt-cache N` (0 = off), `--prompt-cache-every N`, `--turn-token ID`.
**Current limits (v1):** one request at a time, and one conversation cached at a time (switching between two chats
re-reads the other one); greedy decoding (temperature is ignored); images only when set up with them (below); no
video.
---
## Images (vision)
The model has a vision encoder: [`mmproj-Qwen3.8-Flash-Next-BF16.gguf`](https://huggingface.co/ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF)
(0.9 GB, a 27-layer ViT plus the projector into the language model). It is **optional**: say yes when the setup asks
"Images?", or run it again with `--vision gpu` (or `--vision cpu`). The setup downloads the encoder, builds a small
helper (`strata-vision`, from llama.cpp's `mtmd` library) and adds it to your start script. Nothing else changes.
| Encoder on | Time per picture | Cost |
| --- | --- | --- |
| **GPU** (recommended) | **0.1-0.5 s** (up to 1,024 image tokens) | ~1.4 GB of VRAM is kept free for it, so the expert cache is smaller: text output is a few % slower (table below) |
| CPU | 10-30 s (pictures are scaled down to ~300 image tokens) | nothing on the GPU |
A picture becomes up to 1,024 tokens of the context (a 640x480 photo: 300). The same picture sent again, as chat apps
do on every turn, is encoded only once.
### Sending a picture
**Terminal chat:** type `/image <path to a picture>`, press Enter, then type your question.
```
you> /image C:\Users\me\Pictures\receipt.jpg
(picture attached: receipt.jpg - now type your question)
you> What is the total on this receipt?
```
**OpenAI API** (an `image_url` part: a `data:` URL, an `http(s)://` URL or a local file path):
```python
import base64
from openai import OpenAI
client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="none")
img = base64.b64encode(open("photo.jpg", "rb").read()).decode()
r = client.chat.completions.create(model="strata", messages=[{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img}"}},
{"type": "text", "text": "What is in this picture?"}]}])
print(r.choices[0].message.content)
```
**Anthropic API:** an `image` block with a `base64` (or `url`) source, as usual.
JPEG, PNG, BMP, GIF, WebP, TIFF and AVIF work (the last ones are converted to PNG first; agents such as omp send
WebP). Chat apps with image upload work the same way.
### Speed with images on (4K context, measured)
| Model | Prompt tok/s, images off | Prompt tok/s, images on | Output tok/s, images off | Output tok/s, images on |
| --- | ---: | ---: | ---: | ---: |
| **Q2_0** | 541 | 531 | 90.7 | 86.9 |
| **IQ2_XS** | 468 | 458 | 77.1 | 74.0 |
| **IQ3_XXS** | 411 | 401 | 64.7 | 59.5 |
"Off" is the published setup; "on" runs with the encoder loaded on the GPU and its VRAM kept free, so ~1,000 fewer
experts fit in VRAM and the CPU computes a few more per token: 2-8% slower. For text, turning images on changes
nothing else (the output is bit-identical when the VRAM is the same).
A question about a picture (a 640x480 newspaper page = 300 image tokens, a 328-token prompt; the whole request, encoder
on the GPU, measured through the API):
| Model | Answer ("MEN WALK ON MOON") | Picture encoded | Prompt | Output |
| --- | ---: | ---: | ---: | ---: |
| **Q2_0** | 1.9 s | ~0.1 s | 207 tok/s | 65-73 tok/s |
| **IQ2_XS** | 2.3 s | ~0.1 s | 173 tok/s | 50-60 tok/s |
| **IQ3_XXS** | 2.7 s | ~0.1 s | 144 tok/s | 41-50 tok/s |
(Short prompts run below the 4K prompt speed: a 2,048-token chunk is where the prompt path is efficient. Short answers
run below the long-output speed: the first rounds have no draft yet.)
**How it works inside:** the encoder turns the picture into rows of the same width as the model's word embeddings;
Strata puts them where the prompt has `<|image_pad|>` tokens and gives each one its 2-D position (row and column in the
picture; the model uses interleaved M-RoPE). Answers match llama.cpp's multimodal implementation token for token on our
test images.
---
## Troubleshooting
| Symptom | What to do |
| --- | --- |
| `the NVIDIA driver is too old` | Update the driver (NVIDIA App or nvidia.com/drivers), restart, run `START-HERE.bat` again. |
| Python or the build tools could not be installed | Install what it names (links are printed), then run it again. Everything already done is kept. |
| `port 8080 is already in use` | Strata is already running (look for its window), or another program uses the port: `START-HERE.bat --port 8081`. |
| `cudaHostRegister ... out of memory` in the log | Normal on Windows: the engine pins the experts in per-layer slices instead. Only a problem if the load then fails. |
| The first start takes minutes | It is reading 34-43 GB into RAM; the second start is faster while the files are in the OS cache. |
| Slow output, disk light busy | Not enough free RAM: close other programs, or choose Q2_0 / IQ2_XS. |
| `prompt ... exceeds the context` | The request is longer than the context you chose: run setup again with a bigger `--context`. |
| Slower than the tables | The monitor plugged into the GPU and other GPU programs take VRAM from the expert cache; RAM running below its rated speed (enable EXPO/XMP in the BIOS) slows the CPU half. |
| `this server was started without the vision encoder` | The model was set up for text only: run setup again with `--vision gpu`. |
| A picture is refused or `cannot read the image` | The file is not a picture Pillow can open (JPEG, PNG, WebP, GIF, BMP, TIFF, AVIF work). |
| 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 (older engines, mostly with pictures) | Run `START-HERE.bat` once to get engine 0.1.2 or newer. It keeps a real margin of VRAM free: the log says `... MiB of VRAM free with everything loaded` and warns when it is close to 0. |
| Anything else | The engine log is `strata-<model>.log` in this folder. |
---
## How it works
<p align="center"><img src="paper/tiers.svg" width="760" alt="memory tiers"></p>
- **GPU (VRAM):** attention and DeltaNet mixers, the gated-residual weights, routers, shared experts, output head, the MTP
draft layer, the KV cache (from 64K: only its most-read part, the rest streams from RAM), and an **expert cache** that fills the rest of VRAM with the most-used experts (it adapts to
the conversation while you chat).
- **RAM:** all 24,576 experts, pinned. The CPU computes the experts that are not on the GPU **in place**, at the same time
as the GPU works on the cached ones (AVX-512 / AVX2 kernels, ggml's for the i-quants).
- **SSD:** the 28.8 GB n-gram table, read a few rows per token through the OS cache.
- **Speculation:** the model's own MTP layer drafts up to 3 tokens; one pass over all 48 layers checks them. 2.4-3.2
tokens per pass on average.
- **Prompts** are processed in 2,048-token chunks with the experts streamed to the GPU over PCIe.
The full story, with measurements, bottlenecks and what comes next: **[docs/paper/Strata-Paper.pdf](paper/Strata-Paper.pdf)**.
---
## Credits and licenses
- Model: [Qwen/Qwen3.8-Flash-Next](https://huggingface.co/Qwen/Qwen3.8-Flash-Next) by the Qwen team; quantizations:
[ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF](https://huggingface.co/ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF).
Swift 1.5: [ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF](https://huggingface.co/ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF)
by UkisAI. Their licenses apply to the weights.
- [llama.cpp / ggml](https://github.com/ggml-org/llama.cpp) (MIT): the i-quant formats, the GPU dot products and
dequantizers transcribed in `src/kernels/cuda/iq_kernels.cu`, the CPU backend linked for the i-quant experts, the
`mtmd` library behind the image encoder (`tools/vision/`), and `gguf-py` used by the tools. See
`third_party/ggml/LICENSE`.
- Ideas from [Splash](https://github.com/incoai/splash), [ninfer](https://github.com/Neroued/ninfer) and
[HyperQwen](https://github.com/syv-ai/HyperQwen); references in the paper.
+29 -3
View File
@@ -36,6 +36,7 @@
#include "strata/core/weights.hpp"
#include "strata/kernels/gr.hpp"
#include "strata/kernels/kv_stream.hpp"
#include "strata/kernels/ngram.hpp"
#include "strata/kernels/ple.hpp"
@@ -206,10 +207,20 @@ struct QsaState {
int8_t* v_q = nullptr;
uint16_t* k_scale = nullptr;
uint16_t* v_scale = nullptr;
int32_t* page_table = nullptr; ///< (n_pages,) logical page -> physical page
int32_t* page_table = nullptr; ///< (n_pages,) logical page -> physical page (-1: not resident, streamed)
int64_t n_pages = 0;
int64_t max_cells = 0;
/// KV STREAMING (docs/kv-streaming-design.md, `kv_stream.hpp`). `kv_mode` 0: every page in VRAM, identity
/// table (n_slots == n_pages, no host copy). 1: streamed - the authoritative K/V in `host`, `n_slots` pages
/// in VRAM, `map` resolves the selection's blocks on device. 2: a ring (the MTP drafter's window) - `n_slots`
/// pages at `block % n_slots`, with the host copy for a resume. The pool pointers above point at the slots.
int kv_mode = 0;
int64_t n_slots = 0;
strata::kernels::KvHostPools host;
strata::kernels::KvStreamMap map;
int64_t idx_pooled_rows = 0; ///< rows of `idx_pooled` (a ring, which has no indexer, keeps 2)
float* idx_tail = nullptr; ///< (idx_block - 1, idx_dim): the raw tail of the block being filled
float* idx_dead = nullptr; ///< (idx_dim,): the spare slot's key, CONSTANT for the sequence
float* idx_pooled = nullptr; ///< (max_cells/idx_block + 2, idx_dim)
@@ -242,13 +253,28 @@ struct QsaState {
/// Plan v0.3 P7: the RoPE cos/sin table (max_cells x n_rot/2 x 2 floats, 64 MiB at 262K) is identical in every
/// QSA layer. `with_rope = false` sizes a state that borrows it; `share_rope` points `st` at another state's table
/// instead of building a copy (the session builds it once, in the first QSA layer).
uint64_t qsa_state_bytes(const ModelGeometry& g, int64_t max_cells, bool with_rope = true);
uint64_t qsa_state_bytes(const ModelGeometry& g, int64_t max_cells, bool with_rope = true, int64_t ring_cells = 0);
/// KV streaming: keep `cells` cells of each QSA layer in VRAM and the rest in pinned host memory (0: all in VRAM,
/// the default). Set before sizing and initializing the session; a context that fits in `cells` is not streamed.
/// Also puts the MTP drafter's K/V in a ring of its window (`ring_cells` of qsa_state_bytes/init; -1 forces a fully
/// resident state).
void qsa_set_kv_resident(int64_t cells);
int64_t qsa_kv_resident();
/// The fewest resident cells a streamed layer may have: one verify window's selections (8 queries x 2,051 cells
/// in whole blocks) must fit at once, with room to spare.
int64_t qsa_kv_resident_min();
/// Pinned host bytes the streamed states hold (their host copies).
uint64_t qsa_kv_host_bytes();
/// Plan v0.3 P7: store K/V as INT8 with FP16 scales per 64 values (half the VRAM of FP16). Set before sizing and
/// initializing the session; default off until gate G-C accepts it.
void qsa_set_kv_int8(bool enabled);
bool qsa_kv_int8();
uint64_t qsa_state_init(const ModelGeometry& g, int64_t max_cells, void* base, QsaState& st,
const QsaState* share_rope = nullptr);
const QsaState* share_rope = nullptr, int64_t ring_cells = 0);
/// KV streaming: the pools a reader sees (the VRAM slots) and, when streamed, make the selection's blocks resident.
strata::kernels::QsaAttnPools qsa_attn_pools(const QsaState& st);
void qsa_kv_resolve(const QsaState& st, const ModelGeometry& g, const int32_t* ids, const int32_t* steps, int64_t n_q,
int64_t cap, void* stream);
/// Zeroes the pools AND the indexer, so a fresh sequence matches the reference's own `zeros()`. The KV pool
/// matters even for cells that are never attended, because `kv_gather` reads whatever the selection names.
void qsa_state_zero(const QsaState& st, const ModelGeometry& g, void* stream);
+8
View File
@@ -45,7 +45,14 @@ public:
int64_t window = 32768);
/// The prompt's length: prefill() skips the cells the attention window can never reach again.
void set_prompt_len(int64_t n) { prompt_len_ = n; }
/// At most this many drafts per round (below max_t - 1): a window longer than the MTP's comes from elsewhere.
void set_max_drafts(int k) { max_drafts_ = k; }
uint64_t vram_bytes() const { return vram_; }
/// The draft layer's K/V state (read-only: --serve's STRATA_STATE_HASH check hashes it)
const QsaState& kv_state() const { return st_; }
/// KV streaming: refill the ring of the drafter's window from its host copy for a sequence that continues at
/// `upto` (a conversation-cache resume). No-op unless the drafter's K/V is a ring.
void kv_restore(int64_t upto);
/// The main model's embedding and head, and the verify window's final residuals (T rows, hc*n_embd each).
bool bind(const WeightTable& wt, const NativeHead* head, const float* window_R, std::string& err);
@@ -82,6 +89,7 @@ private:
const NativeHead* head_ = nullptr;
const float* window_R_ = nullptr;
int max_t_ = 0;
int max_drafts_ = 1 << 30;
int64_t n_vocab_ = 0;
uint64_t vram_ = 0;
cudaStream_t cs_ = nullptr;
+17
View File
@@ -23,7 +23,10 @@
#include "strata/kernels/cpu/native_expert.hpp"
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <mutex>
#include <thread>
#include <vector>
@@ -136,10 +139,19 @@ public:
repark = ms_repark_;
}
/// **A PARKED WORKER SPINS FOR THIS LONG, THEN SLEEPS.** The park is a `_mm_pause` spin because a layer's
/// batches are microseconds apart and a wake-up from the OS costs more than that. But a spin that never ends
/// keeps every worker's core at 100% while the engine waits for a request - issue #4, "CPU 50% even when
/// doing nothing" (7 of the 5700X's 16 threads). Inside a request the gaps between batches are far below
/// this, so the token path still never sleeps; between requests the workers block on `sleep_cv_`.
static constexpr std::chrono::milliseconds kSpinBeforeSleep{20};
private:
void worker(int id);
void drain(int id, ExpertScratch& scratch);
void run_phase(int mode, int n_tasks);
/// Bump `epoch_`, and wake the workers that went to sleep. Every publish goes through here.
void publish();
int n_ = 0;
bool host_works_ = true;
@@ -169,6 +181,11 @@ private:
alignas(64) std::atomic<uint32_t> parked_{0};
alignas(64) std::atomic<uint32_t> epoch_{0};
alignas(64) std::atomic<bool> stop_{false};
// The sleep after `kSpinBeforeSleep`. `sleepers_` is how `publish` knows whether anyone needs waking, so the
// token path pays one uncontended load per publish and never takes the mutex while the workers spin.
alignas(64) std::atomic<uint32_t> sleepers_{0};
std::mutex sleep_mu_;
std::condition_variable sleep_cv_;
std::vector<std::thread> threads_;
std::vector<ExpertScratch> scratch_; // one per worker: no allocation, no false sharing of the hot data
// run_split state: mode 0 = whole experts, 1 = gate/up row parts, 2 = down row parts
+5 -2
View File
@@ -12,6 +12,7 @@
// the FP16 pools: codes `[page][kv_head][page_size][head_dim]`, scales `[page][kv_head][page_size][head_dim/64]`.
#pragma once
#include "strata/kernels/kv_stream.hpp"
#include "strata/kernels/qsa.hpp"
#include <cstdint>
@@ -25,9 +26,11 @@ inline uint64_t kv_q8_bytes_per_cell(const QsaShapes& s) {
return (uint64_t) s.n_head_kv * s.head_dim * 2 + (uint64_t) s.n_head_kv * (s.head_dim / KV_Q8_GROUP) * 2 * 2;
}
/// Append the cell at step[kStepPos] (graph-capturable: position and page come from device memory).
/// Append the cell at step[kStepPos] (graph-capturable: position and page come from device memory). With a host
/// copy (KV streaming, `kv_stream.hpp`) the cell is written there too, and to VRAM only if its block is resident.
void kv_append_q8_step(int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, const int32_t* page_table,
const int32_t* step, const float* kcur, const float* vcur, const QsaShapes& s, void* stream);
const int32_t* step, const float* kcur, const float* vcur, const QsaShapes& s, void* stream,
const KvHostPools* host = nullptr);
/// Gather step[kStepWidth] cells named by `ids` into FP16 scratch `[id][kv_head][head_dim]`; the grid is sized by
/// `max_ids` (capacity), the kernel reads the real count from `step`.
+92
View File
@@ -0,0 +1,92 @@
// include/strata/kernels/kv_stream.hpp - KV streaming for the QSA layers (docs/kv-streaming-design.md).
//
// A streamed QSA layer keeps its AUTHORITATIVE K/V in pinned, device-mapped host memory, laid out exactly as a
// fully resident pool with the identity page table would be. VRAM holds a pool of `n_slots` pages (one page =
// `page_size` = 4 cells = one indexer block, all KV heads, K and V: 4,224 B in int8), and the page table becomes a
// RESIDENCY MAP: `page_table[block]` is the slot holding that block, or -1.
//
// Every reader (`qsa_decode_attn_*`, `kv_gather_*`) is unchanged: it goes through the page table as before. What
// changes is that after the selection and before the attention, `kv_stream_resolve` makes every block the
// selection names resident - on device, inside the captured graphs:
// 1. one block of 1,024 threads walks the selected ids: a resident block gets this call's epoch and its clock
// reference bit, a missing block is claimed once (page_table -1 -> -2) into a miss list;
// 2. the same block picks one victim slot per miss with a CLOCK sweep (second chance) that never takes a slot
// this call uses, re-points the page table (old block -> -1, new block -> slot);
// 3. a copy kernel reads the missed blocks from the host copy (zero-copy, over PCIe) into their slots.
// Writers (`kv_append_*`) write the host copy always, and the slot only if the block is resident, so the slots
// never go stale and nothing needs invalidating - not after a rejected draft, not on a conversation-cache resume.
//
// A RING (the MTP drafter, which only ever reads its last `window` cells) uses the same host copy with a static
// table `block -> block % n_slots` and no resolve: the ring is refilled from the host copy on a resume.
#pragma once
#include "strata/kernels/qsa.hpp"
#include "strata/kernels/qsa_decode_attn.hpp"
#include <cstdint>
namespace strata::kernels {
/// The host copy of one layer's K/V: device-mapped pointers (UVA) into pinned memory, identity layout
/// `[block][kv_head][page_size][head_dim]`. All null when the layer is fully resident.
struct KvHostPools {
uint16_t* k_pool = nullptr; ///< fp16 mode
uint16_t* v_pool = nullptr;
int8_t* k_q = nullptr; ///< int8 mode: codes
int8_t* v_q = nullptr;
uint16_t* k_scale = nullptr; ///< int8 mode: fp16 scale per 64 values
uint16_t* v_scale = nullptr;
bool present() const { return k_pool != nullptr || k_q != nullptr; }
};
/// The residency map of a streamed layer, all device memory at fixed addresses (the graphs bake them in).
struct KvStreamMap {
int32_t* page_table = nullptr; ///< (n_blocks,) block -> slot, or -1
int32_t* slot_block = nullptr; ///< (n_slots,) slot -> block, or -1
int32_t* slot_stamp = nullptr; ///< (n_slots,) the resolve epoch that last used the slot
int32_t* slot_ref = nullptr; ///< (n_slots,) clock reference bit
int32_t* ctl = nullptr; ///< kKvCtlInts: epoch, hand, misses of the last call, overflow, u64 counters
int32_t* miss_block = nullptr; ///< (n_slots,)
int32_t* miss_slot = nullptr; ///< (n_slots,)
int64_t n_blocks = 0;
int64_t n_slots = 0;
};
inline constexpr int kKvCtlInts = 16; ///< [0] epoch [1] hand [2] misses [3] overflow; u64 at [4] misses, [6] lookups, [8] calls
/// Bytes of the map's arrays besides the page table (which the state already has): 5 per-slot ints + ctl.
inline uint64_t kv_stream_map_bytes(int64_t n_slots) { return (uint64_t) n_slots * 4 * 5 + kKvCtlInts * 4; }
/// Every block evicted: page table all -1, slots free, counters zero.
void kv_stream_reset(const KvStreamMap& m, void* stream);
/// Make every block named by the selections of `n_q` queries resident (ids [n_q][cap], width from
/// steps[q * kStepCount + kStepWidth]). Capturable. `n_slots` must hold the distinct blocks of one call
/// (n_q x (cap / page_size + 2)); a call that cannot sets ctl[3] (see `kv_stream_counters`).
void kv_stream_resolve(const KvStreamMap& m, const QsaAttnPools& slots, const KvHostPools& host, bool int8,
const int32_t* ids, const int32_t* steps, int64_t n_q, int64_t cap, const QsaShapes& s,
void* stream);
/// The static ring table `block -> block % n_slots`.
void kv_ring_table(int32_t* page_table, int64_t n_blocks, int64_t n_slots, void* stream);
/// Copy blocks [b0, b1) of the host copy into the pool pages `page_table` names (host-side table `phys(b)`
/// computed as `b % n_slots`: the ring restore). Not capturable.
void kv_ring_restore(const QsaAttnPools& slots, const KvHostPools& host, bool int8, int64_t b0, int64_t b1,
int64_t n_slots, const QsaShapes& s, void* stream);
/// Copy the first `n_blocks` blocks of the host copy into a fully resident (identity-layout) pool: the prompt
/// path's staging of one layer. Not capturable (DMA).
void kv_stage_from_host(const QsaAttnPools& stage, const KvHostPools& host, bool int8, int64_t n_blocks,
const QsaShapes& s, void* stream);
struct KvStreamCounters {
uint64_t misses = 0, lookups = 0, calls = 0;
bool overflow = false;
};
/// Synchronous read of the counters (debug and the end-of-request summary).
KvStreamCounters kv_stream_counters(const KvStreamMap& m);
/// Bytes of one block (page) of K and V together.
uint64_t kv_block_bytes(const QsaShapes& s, bool int8);
} // namespace strata::kernels
+9 -3
View File
@@ -110,7 +110,9 @@ inline QsaShapes qsa_real_shapes() {
QsaShapes s;
s.n_head = 24; s.n_head_kv = 2; s.head_dim = 256; s.n_rot = 64;
s.idx_n_head = 4; s.idx_dim = 128; s.idx_block = 4; s.idx_top_k = 2048;
s.page_size = 512;
// one page = one indexer block (4 cells): the granule KV streaming keeps resident (kv_stream.hpp). The readers
// resolve a row per cell anyway, so the page size costs them nothing (measured: identical output and speed).
s.page_size = 4;
return s;
}
@@ -353,8 +355,12 @@ void qsa_attend_step(const float* q, const uint16_t* k_scratch, const uint16_t*
const int32_t* step, int64_t max_ids, const QsaShapes& s, float* attn, float* weights,
void* stream);
/// The cell's position comes from `step`, so this is the form a graph may contain.
struct KvHostPools; // kv_stream.hpp
/// The cell's position comes from `step`, so this is the form a graph may contain. With a host copy (KV streaming)
/// the cell is written there too, and to VRAM only if its block is resident.
void kv_append_step(uint16_t* k_pool, uint16_t* v_pool, const int32_t* page_table, const int32_t* step,
const float* kcur, const float* vcur, const QsaShapes& s, void* stream);
const float* kcur, const float* vcur, const QsaShapes& s, void* stream,
const KvHostPools* host = nullptr);
} // namespace strata::kernels
+6 -1
View File
@@ -5,6 +5,8 @@
// the kernel walks the chunk's tokens in order inside one launch.
#pragma once
#include "strata/kernels/kv_stream.hpp"
#include <cstdint>
namespace strata::prefill {
@@ -63,11 +65,14 @@ void gate_attn(const float* attn, const float* q_full, uint16_t* out16, int64_t
/// K, V: [T, 2, 256].
void kv_append(const float* K, const float* V, int64_t T, int64_t pos0, const int32_t* page_table, int64_t page_size,
uint16_t* k_pool, uint16_t* v_pool, int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale,
void* stream);
void* stream, const strata::kernels::KvHostPools* host = nullptr,
const strata::kernels::KvHostPools* stage = nullptr);
/// fp32 -> fp16 bits and fp32 -> bf16, n elements (the two activation images of the prompt GEMMs).
void to_f16(const float* x, uint16_t* y, int64_t n, void* stream);
void to_bf16(const float* x, uint16_t* y, int64_t n, void* stream);
/// y = fp32(fp16(x)): what an FP16 store of x would read back (the FP16 indexer-key experiment)
void round_f16(const float* x, float* y, int64_t n, void* stream);
/// Expert blob -> FP16 (Q2_0 values are exact in FP16).
void blob_dequant_f16(const uint8_t* blob, uint16_t* gu16, uint16_t* down16, void* stream);
+3
View File
@@ -62,6 +62,9 @@ public:
/// K/V from them. The prefill stream is synchronized before the call.
std::function<bool(const float* R_rows, int64_t T, int64_t pos0, std::string& err)> on_chunk;
/// Checked before every chunk: true stops the prompt early (`run` returns false with err "cancelled").
std::function<bool()> should_stop;
/// The vision path: HOST rows (n_embd floats) indexed by absolute position, read in place of the token
/// embedding where non-null (an image's <|image_pad|> cells). Null (default): every position embeds its token.
const float* const* embd_rows = nullptr;
+137 -2
View File
@@ -270,11 +270,127 @@ class OutputParser:
"""Incremental parser of the model's text. Feed deltas; get events. A tag split across deltas is held back
until it is complete, so clients never see `<tool_` or `</thi`."""
def __init__(self, thinking: bool = True, tools: list[dict] | None = None):
def __init__(self, thinking: bool = True, tools: list[dict] | None = None, stream_tools: bool = False):
self.state = "reasoning" if thinking else "content"
self.buf = ""
self.lead = False
self.schemas = {t.get("name"): t for t in tools or []}
# stream_tools: a tool call is also reported while it is being written - "tool_start" (its name and id) as
# soon as the name is known, then "tool_args" pieces of its JSON arguments (string parameters character by
# character; other types whole, once complete) - before the final "tool_call". Without it, a client sees
# nothing until the call is complete, which for a large file write can be many minutes.
self.stream_tools = stream_tools
self._reset_scan()
def _reset_scan(self):
self.sp = 0 # how much of self.buf (the call body) the scanner has consumed
self.ss = "name" # name -> between -> str|raw -> ... -> done
self.scall = None # the ToolCall being streamed (its id is reused by the final event)
self.sfirst = True
self.sval_started = False
self.sdeclared = {}
def _scan(self) -> list[Event]:
"""Advance the streaming view of the call body in self.buf (see stream_tools)."""
out = []
def args(s):
if s:
out.append(Event("tool_args", s, call=self.scall))
while True:
rest = self.buf[self.sp:]
if self.ss == "name":
a = rest.find("<function=")
b = rest.find(">", a + 10) if a >= 0 else -1
if b < 0:
return out
name = rest[a + 10:b]
self.scall = ToolCall(name=name, arguments={})
props = ((self.schemas.get(name) or {}).get("parameters") or {}).get("properties") or {}
self.sdeclared = {k: (v or {}).get("type") for k, v in props.items()}
out.append(Event("tool_start", call=self.scall))
args("{")
self.sp += b + 1
self.ss = "between"
elif self.ss == "between":
stripped = rest.lstrip()
self.sp += len(rest) - len(stripped)
if stripped.startswith("<parameter="):
b = stripped.find(">")
if b < 0:
return out
pname = stripped[11:b]
args(("" if self.sfirst else ",") + json.dumps(pname) + ":")
self.sfirst = False
self.sp += b + 1
if self.sdeclared.get(pname) == "string":
args('"')
self.ss, self.sval_started = "str", False
else:
self.ss = "raw"
elif stripped.startswith("</function>"):
args("}")
self.sp += len("</function>")
self.ss = "done"
else:
return out # a tag still arriving (or trailing text): wait
elif self.ss == "str":
if not self.sval_started:
if not rest:
return out
if rest[0] == "\n": # the value's leading newline is not part of it
self.sp += 1
rest = rest[1:]
self.sval_started = True
end = rest.find("</parameter>")
if end >= 0:
value = rest[:end]
if value.endswith("\n"):
value = value[:-1]
args(json.dumps(value)[1:-1] + '"')
self.sp += end + len("</parameter>")
self.ss = "between"
continue
safe = len(rest) - self._hold(rest, ("</parameter>",))
if safe > 0 and rest[safe - 1] == "\n": # may be the trailing newline before </parameter>
safe -= 1
if safe > 0:
args(json.dumps(rest[:safe])[1:-1])
self.sp += safe
return out
elif self.ss == "raw":
end = rest.find("</parameter>")
if end < 0:
return out
value = rest[:end]
if value.startswith("\n"):
value = value[1:]
if value.endswith("\n"):
value = value[:-1]
try:
v = json.loads(value)
except ValueError:
v = value
args(json.dumps(v, ensure_ascii=False))
self.sp += end + len("</parameter>")
self.ss = "between"
else:
return out
def _close_scan(self) -> list[Event]:
"""A streamed call that ended without a clean </function>: close its JSON so clients can still parse it."""
out = []
if self.scall is None or self.ss == "done":
return out
tail = ""
if self.ss == "str":
tail += '"'
elif self.ss == "raw":
tail += json.dumps(self.buf[self.sp:].strip("\n"))
tail += "}"
out.append(Event("tool_args", tail, call=self.scall))
self.ss = "done"
return out
def _hold(self, text: str, tags: tuple[str, ...]) -> int:
"""Length of the longest suffix of `text` that is a proper prefix of one of `tags`."""
@@ -325,17 +441,36 @@ class OutputParser:
self.state = "call"
else:
i = self.buf.find(CALL_END)
if self.stream_tools:
if i >= 0:
whole, self.buf = self.buf, self.buf[:i] # scan only the body
out += self._scan()
out += self._close_scan()
self.buf = whole
else:
out += self._scan()
if i < 0:
return out
body = self.buf[:i]
self.buf = self.buf[i + len(CALL_END):]
name = body.strip()[len("<function="):].split(">", 1)[0]
out.append(Event("tool_call", call=parse_tool_call(body, self.schemas.get(name))))
call = parse_tool_call(body, self.schemas.get(name))
if self.scall is not None:
call.id = self.scall.id
out.append(Event("tool_call", call=call))
self._reset_scan()
self.state, self.lead = "content", True
def finish(self) -> list[Event]:
"""End of generation: flush whatever is held (an unterminated tool call is returned as content)."""
out = []
if self.state == "call" and self.stream_tools and self.scall is not None:
out += self._scan() # the output ended inside a call that was already announced
out += self._close_scan()
out.append(Event("tool_call", call=self.scall))
self.buf = ""
self._reset_scan()
return out
if self.buf:
kind = {"reasoning": "reasoning", "content": "content"}.get(self.state, "content")
text = self.buf if self.state != "call" else CALL_START + self.buf
+272 -39
View File
@@ -6,8 +6,9 @@
Endpoints: POST /v1/chat/completions (OpenAI, stream and non-stream), POST /v1/messages (Anthropic, stream and
non-stream), GET /v1/models, GET /health. One sequence at a time behind a FIFO (plan: one resident sequence).
Images (optional, when the config has a "vision" entry): OpenAI image_url parts and Anthropic image blocks (base64
data, http(s) URLs or local file paths; JPEG/PNG/BMP/GIF) go through `strata-vision` (the model's mmproj file) and
reach the engine as embeddings (`GENI`).
data, http(s) URLs or local file paths) go through `strata-vision` (the model's mmproj file) and reach the engine as
embeddings (`GENI`). JPEG/PNG/BMP/GIF go straight in; WebP, TIFF, AVIF, ... (agents like omp send WebP) are
converted to PNG first with Pillow.
Requests whose prompt plus max tokens exceed the engine's context are REJECTED with 400, never truncated.
The engine boundary is `Engine.generate(prompt_ids, max_new, sampling, cancel) -> iterator of token ids`.
@@ -21,6 +22,7 @@ import base64
import hashlib
import json
import os
import queue
import subprocess
import sys
import tempfile
@@ -77,38 +79,79 @@ class StrataEngine:
self.proc = subprocess.Popen([exe, "--serve", *args], cwd=cwd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
stderr=self.log, text=True, encoding="utf-8", bufsize=1, env=env)
self.max_context = 0
self.can_stop = False # the engine honours a STOP line mid-request (READY <ctx> stop)
self.last = {}
for line in self.proc.stdout:
if line.startswith("READY"):
self.max_context = int(line.split()[1])
f = line.split()
self.max_context = int(f[1])
self.can_stop = "stop" in f[2:]
break
if self.max_context <= 0:
raise RuntimeError("the engine exited before it was ready" + (f" (see {log})" if log else ""))
# the engine's stdout on a thread, so a request can wait with a timeout (heartbeats, cancel checks)
self.lines: queue.Queue = queue.Queue()
threading.Thread(target=self._pump, daemon=True).start()
def _pump(self):
for line in self.proc.stdout:
self.lines.put(line)
self.lines.put(None)
def _parse_done(self, line):
f = line.split()
self.last = {"generated": int(f[1]), "prompt_tokens": int(f[2]), "prompt_ms": float(f[3]),
"decode_ms": float(f[4]), "finish": f[5]}
def generate(self, ids, max_new, sampling, cancel, embeddings=None):
"""Yields token ids, and None as a heartbeat every 10 s while the engine is quiet (reading a long prompt):
the HTTP layer turns it into an SSE comment, which keeps clients' watchdogs calm and notices a client that
has gone. A consumer that stops early (or `cancel`) makes the engine STOP, so it does not run to max_new."""
head = f"GENI {int(max_new)} {embeddings}" if embeddings else f"GEN {int(max_new)}"
self.proc.stdin.write(f"{head} {','.join(str(int(t)) for t in ids)}\n")
self.proc.stdin.flush()
done = False
try:
for line in self.proc.stdout:
while True:
try:
line = self.lines.get(timeout=10)
except queue.Empty:
if cancel.is_set():
return
yield None
continue
if line is None:
done = True
raise RuntimeError("the engine process ended")
if line.startswith("T "):
if not cancel.is_set():
yield int(line[2:])
if cancel.is_set():
return
yield int(line[2:])
elif line.startswith("PP "): # prompt progress, one per chunk: also a heartbeat (the
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("DONE"):
f = line.split()
self.last = {"generated": int(f[1]), "prompt_tokens": int(f[2]), "prompt_ms": float(f[3]),
"decode_ms": float(f[4]), "finish": f[5]}
self._parse_done(line)
done = True
return
elif line.startswith("ERR"):
done = True
raise ValueError(line[4:].strip())
raise RuntimeError("the engine process ended")
finally:
if not done: # the consumer stopped early: drain to DONE
for line in self.proc.stdout:
if line.startswith("DONE") or line.startswith("ERR"):
if not done: # the consumer stopped early: stop the engine, drain to DONE
if self.can_stop:
try:
self.proc.stdin.write("STOP\n")
self.proc.stdin.flush()
except OSError:
pass
while True:
line = self.lines.get()
if line is None or line.startswith("ERR"):
break
if line.startswith("DONE"):
self._parse_done(line)
break
def close(self):
@@ -155,9 +198,37 @@ class Vision:
return Path(path).read_bytes()
raise ValueError("an image must be a data: URL, an http(s) URL or a local file path")
@staticmethod
def normalize(data: bytes) -> bytes:
"""The formats strata-vision's decoder (stb_image) reads pass through; anything else is converted to PNG."""
if data[:3] == b"\xff\xd8\xff" or data[:8] == b"\x89PNG\r\n\x1a\n" or data[:2] == b"BM" or \
data[:6] in (b"GIF87a", b"GIF89a"):
return data
try:
import io
from PIL import Image
except ImportError:
raise ValueError("this image format needs Pillow (python -m pip install pillow); JPEG, PNG, BMP and "
"GIF work without it") from None
try:
im = Image.open(io.BytesIO(data))
im.load()
except Exception as e:
raise ValueError(f"the image could not be read ({e})") from None
if im.mode in ("RGBA", "LA", "P") and "transparency" in im.info or im.mode in ("RGBA", "LA"):
im = im.convert("RGBA")
bg = Image.new("RGB", im.size, (255, 255, 255)) # transparent areas become white, not black
bg.paste(im, mask=im.split()[-1])
im = bg
elif im.mode != "RGB":
im = im.convert("RGB")
out = io.BytesIO()
im.save(out, format="PNG")
return out.getvalue()
def encode(self, source: str) -> tuple[Path, int]:
"""-> (embeddings file, number of image tokens)."""
data = self.load(source)
data = self.normalize(self.load(source))
key = hashlib.sha256(data).hexdigest()[:32]
with self.lock:
if key in self.cache:
@@ -245,6 +316,8 @@ class Service:
self.fifo = threading.Lock()
self.embeddings = threading.local() # the current request's image embeddings file (GENI)
self.api_key = "" # when set, /v1/* needs it (Bearer or x-api-key)
self.status = {"busy": False, "queued": 0} # GET /status: what the model is doing right now
self.status_lock = threading.Lock()
self.stop_ids = set(tokenizer.encode(IM_END, parse_special=True) +
tokenizer.encode("<|endoftext|>", parse_special=True))
@@ -258,7 +331,12 @@ class Service:
raise ValueError("this server was started without the vision encoder (run setup again and choose "
"'vision'), so it cannot read images")
pad = self.tok.encode(IMAGE_PAD, parse_special=True)[0]
encoded = [self.vision.encode(src) for src in images]
# Encode only while the engine is idle: the engine and the image encoder (a separate process) must not
# run on the GPU at the same time - an encode during a running request left that request stuck at
# "reading the prompt" with CPU and GPU busy, for good (reproduced). So encoding takes its turn in the
# same FIFO as the requests.
with self.fifo:
encoded = [self.vision.encode(src) for src in images]
out, k = [], 0
for t in ids: # one <|image_pad|> per image -> one per image token
if t == pad and k < len(encoded):
@@ -279,29 +357,116 @@ class Service:
f"({self.engine.max_context}); requests are never truncated")
return ids, kwargs.get("enable_thinking", True) is not False
def _note(self, n, evs):
with self.status_lock:
s = self.status
s["generated"] = n
if s.get("first_token") is None:
s["first_token"] = time.time()
for ev in evs:
if ev.kind == "reasoning":
s["phase"] = "thinking"
elif ev.kind == "content":
s["phase"] = "answering"
elif ev.kind == "tool_start":
s["phase"], s["tool"] = f"writing a tool call: {ev.call.name}", ev.call.name
elif ev.kind == "tool_call":
s["phase"] = "tool call complete"
s["tail"] = (s["tail"] + (ev.text or ""))[-600:]
def _progress(self, last_print, every=1.0):
"""A progress line in the server window every `every` seconds while a request runs."""
now = time.time()
if now - last_print < every:
return last_print
with self.status_lock:
s = dict(self.status)
el = now - s.get("started", now)
if s.get("first_token") is None:
print(f"[strata] reading the prompt: {s.get('prompt_tokens', 0)} tokens, {el:.0f} s so far", flush=True)
else:
rate = s["generated"] / max(1e-6, now - s["first_token"])
print(f"[strata] {s['phase']}: {s['generated']} of max {s.get('max_tokens')} tokens, {rate:.1f} tok/s, "
f"{el:.0f} s", flush=True)
return now
def run(self, ids, thinking, tools, max_new, sampling, cancel) -> Iterator[tuple[str, object]]:
"""Yields ("event", Event) as text arrives, then ("done", {"finish": .., "completion_tokens": ..})."""
parser, detok, n, finish = OutputParser(thinking=thinking, tools=tools), Detokenizer(self.tok), 0, "length"
parser = OutputParser(thinking=thinking, tools=tools, stream_tools=True)
detok, n, finish = Detokenizer(self.tok), 0, "length"
raw_ids = [] # every generated id (STRATA_DEBUG: dump raw model text)
emb = getattr(self.embeddings, "path", None)
with self.status_lock:
self.status["queued"] += 1
try:
with self.fifo:
with self.status_lock:
self.status["queued"] -= 1
self.status.update(busy=True, phase="reading the prompt", prompt_tokens=len(ids), generated=0,
started=time.time(), first_token=None, tool=None, tail="", max_tokens=max_new)
last_print = time.time()
gen = self.engine.generate(ids, max_new, sampling, cancel, embeddings=emb) if emb else \
self.engine.generate(ids, max_new, sampling, cancel)
for t in gen:
n += 1
if t in self.stop_ids:
finish = "stop"
break
for ev in parser.feed(detok.push(t)):
yield "event", ev
try:
for t in gen:
if t is None: # heartbeat while the engine is quiet
last_print = self._progress(last_print)
yield "ping", None
continue
n += 1
if t in self.stop_ids:
finish = "stop"
raw_ids.append(t)
break
raw_ids.append(t)
evs = parser.feed(detok.push(t))
self._note(n, evs)
last_print = self._progress(last_print)
for ev in evs:
yield "event", ev
if cancel.is_set():
finish = "cancel"
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
except GeneratorExit: # the client disconnected mid-stream
finish = "disconnect"
raise
finally:
if emb:
Path(emb).unlink(missing_ok=True)
with self.status_lock:
if self.status.get("busy"):
now = time.time()
el = now - self.status.get("started", now)
ft = self.status.get("first_token")
rate = n / max(1e-6, now - ft) if ft else 0.0
print(f"[strata] done: {n} tokens in {el:.0f} s ({rate:.1f} tok/s) "
f"({finish}, cancel={cancel.is_set()})", flush=True)
if os.environ.get("STRATA_DEBUG") and raw_ids:
print(f"[strata] raw: {self.tok.decode(raw_ids)!r}", flush=True)
self.status["busy"] = False
for ev in parser.finish():
yield "event", ev
yield "done", {"finish": finish, "completion_tokens": n}
def _debug_req(api, req, messages, tools, max_new, thinking, prompt_tokens):
"""One compact line per request while diagnosing blank/empty turns. Set STRATA_DEBUG=1 to enable."""
if not os.environ.get("STRATA_DEBUG"):
return
last = messages[-1] if messages else {}
body = last.get("content")
if isinstance(body, list):
body = " ".join(p.get("text", "") for p in body if isinstance(p, dict))
preview = (str(body or "")[:80]).replace("\n", " ")
print(f"[strata] req {api}: msgs={len(messages)} tools={len(tools or [])} "
f"max_tokens_raw={req.get('max_tokens')!r}/{req.get('max_completion_tokens')!r} "
f"max_new={max_new} thinking={thinking} stream={bool(req.get('stream'))} "
f"prompt_tokens={prompt_tokens} last={last.get('role')!r}:{preview!r}", flush=True)
# ------------------------------------------------------------------------------------------------ OpenAI
def openai_chunks(svc: Service, req: dict, ids, thinking, tools, max_new, cancel):
cid, created = "chatcmpl-" + uuid.uuid4().hex[:24], int(time.time())
@@ -312,20 +477,32 @@ def openai_chunks(svc: Service, req: dict, ids, thinking, tools, max_new, cancel
yield chunk({"role": "assistant", "content": ""})
calls = 0
streamed = {} # tool call id -> index, for calls sent piece by piece
for kind, x in svc.run(ids, thinking, tools, max_new, req, cancel):
if kind == "event":
if kind == "ping":
yield None
elif kind == "event":
ev: Event = x
if ev.kind == "reasoning" and ev.text:
yield chunk({"reasoning_content": ev.text})
elif ev.kind == "content" and ev.text:
yield chunk({"content": ev.text})
elif ev.kind == "tool_start":
streamed[ev.call.id] = calls
calls += 1
yield chunk({"tool_calls": [{"index": streamed[ev.call.id], "id": ev.call.id, "type": "function",
"function": {"name": ev.call.name, "arguments": ""}}]})
elif ev.kind == "tool_args":
yield chunk({"tool_calls": [{"index": streamed[ev.call.id], "function": {"arguments": ev.text}}]})
elif ev.kind == "tool_call" and ev.call.id in streamed:
continue
elif ev.kind == "tool_call":
yield chunk({"tool_calls": [{"index": calls, "id": ev.call.id, "type": "function",
"function": {"name": ev.call.name,
"arguments": json.dumps(ev.call.arguments, ensure_ascii=False)}}]})
calls += 1
else:
finish = "tool_calls" if calls and x["finish"] == "stop" else x["finish"]
finish = "tool_calls" if calls and x["finish"] == "stop" else {"cancel": "stop"}.get(x["finish"], x["finish"])
last = chunk({}, finish)
last["usage"] = {"prompt_tokens": len(ids), "completion_tokens": x["completion_tokens"],
"total_tokens": len(ids) + x["completion_tokens"]}
@@ -333,13 +510,22 @@ def openai_chunks(svc: Service, req: dict, ids, thinking, tools, max_new, cancel
def openai_collect(chunks) -> dict:
content, reasoning, calls, last = [], [], [], None
content, reasoning, by_index, last = [], [], {}, None
for c in chunks:
if c is None: # a heartbeat
continue
d = c["choices"][0]["delta"]
content.append(d.get("content") or "")
reasoning.append(d.get("reasoning_content") or "")
calls += [{k: v for k, v in tc.items() if k != "index"} for tc in d.get("tool_calls") or []]
for tc in d.get("tool_calls") or []: # streamed calls arrive in pieces: merge them by index
cur = by_index.setdefault(tc.get("index", len(by_index)), {"id": None, "type": "function",
"function": {"name": "", "arguments": ""}})
cur["id"] = tc.get("id") or cur["id"]
fn = tc.get("function") or {}
cur["function"]["name"] += fn.get("name") or ""
cur["function"]["arguments"] += fn.get("arguments") or ""
last = c
calls = [by_index[i] for i in sorted(by_index)]
msg = {"role": "assistant", "content": "".join(content) or None}
if "".join(reasoning):
msg["reasoning_content"] = "".join(reasoning)
@@ -361,12 +547,25 @@ def anthropic_events(svc: Service, req: dict, ids, thinking, tools, max_new, can
def close():
return ("content_block_stop", {"type": "content_block_stop", "index": index})
streamed = set()
for kind, x in svc.run(ids, thinking, tools, max_new, req, cancel):
if kind == "ping":
yield None
continue
if kind == "event":
ev: Event = x
want = {"reasoning": "thinking", "content": "text", "tool_call": "tool_use"}[ev.kind]
if ev.kind != "tool_call" and not ev.text:
if ev.kind == "tool_args":
yield "content_block_delta", {"type": "content_block_delta", "index": index,
"delta": {"type": "input_json_delta", "partial_json": ev.text}}
continue
if ev.kind == "tool_call" and ev.call.id in streamed:
continue
want = {"reasoning": "thinking", "content": "text", "tool_call": "tool_use", "tool_start": "tool_use"}[ev.kind]
if ev.kind not in ("tool_call", "tool_start") and not ev.text:
continue
if ev.kind == "tool_start":
streamed.add(ev.call.id)
used_tool = True
if open_kind != want or want == "tool_use":
if open_kind is not None:
yield close()
@@ -383,6 +582,8 @@ def anthropic_events(svc: Service, req: dict, ids, thinking, tools, max_new, can
elif want == "text":
yield "content_block_delta", {"type": "content_block_delta", "index": index,
"delta": {"type": "text_delta", "text": ev.text}}
elif ev.kind == "tool_start":
pass # its input follows as tool_args pieces
else:
used_tool = True
yield "content_block_delta", {"type": "content_block_delta", "index": index, "delta": {
@@ -391,7 +592,7 @@ def anthropic_events(svc: Service, req: dict, ids, thinking, tools, max_new, can
if open_kind is not None:
yield close()
stop = "tool_use" if used_tool and x["finish"] == "stop" else \
{"stop": "end_turn", "length": "max_tokens"}[x["finish"]]
{"stop": "end_turn", "length": "max_tokens", "cancel": "end_turn"}[x["finish"]]
yield "message_delta", {"type": "message_delta", "delta": {"stop_reason": stop, "stop_sequence": None},
"usage": {"output_tokens": x["completion_tokens"]}}
yield "message_stop", {"type": "message_stop"}
@@ -399,7 +600,10 @@ def anthropic_events(svc: Service, req: dict, ids, thinking, tools, max_new, can
def anthropic_collect(events) -> dict:
msg, blocks = None, []
for name, e in events:
for item in events:
if item is None: # a heartbeat
continue
name, e = item
if name == "message_start":
msg = e["message"]
elif name == "content_block_start":
@@ -410,8 +614,11 @@ def anthropic_collect(events) -> dict:
b["text"] += d["text"]
elif d["type"] == "thinking_delta":
b["thinking"] += d["thinking"]
else:
b["input"] = json.loads(d["partial_json"])
else: # input_json_delta pieces: parsed when complete
b["_json"] = b.get("_json", "") + d["partial_json"]
elif name == "content_block_stop" and blocks and "_json" in blocks[-1]:
b = blocks[-1]
b["input"] = json.loads(b.pop("_json") or "{}")
elif name == "message_delta":
msg["stop_reason"] = e["delta"]["stop_reason"]
msg["usage"]["output_tokens"] = e["usage"]["output_tokens"]
@@ -457,6 +664,17 @@ def make_handler(svc: Service):
elif path == "/health":
self._json(200, {"status": "ok", "max_context": svc.engine.max_context, "model": svc.model,
"images": svc.vision is not None, "api_key": bool(svc.api_key)})
elif path == "/status":
with svc.status_lock:
s = dict(svc.status)
now = time.time()
if s.get("busy"):
s["elapsed_s"] = round(now - s["started"], 1)
if s.get("first_token"):
s["tokens_per_s"] = round(s["generated"] / max(1e-6, now - s["first_token"]), 1)
for k in ("started", "first_token"):
s.pop(k, None)
self._json(200, s)
elif path == "/v1/models":
if self._authorized():
self._json(200, {"object": "list", "data": [{"id": svc.model, "object": "model"}]})
@@ -486,7 +704,10 @@ def make_handler(svc: Service):
def _openai(self, req):
messages, tools, kw = openai_to_messages(req)
max_new = int(req.get("max_completion_tokens") or req.get("max_tokens") or 1024)
if max_new <= 0: # a non-positive budget (some clients send -1) means "unset"
max_new = 1024
ids, thinking = svc.prepare(messages, tools, kw, max_new)
_debug_req("openai", req, messages, tools, max_new, thinking, len(ids))
cancel = threading.Event()
chunks = openai_chunks(svc, req, ids, thinking, tools, max_new, cancel)
if not req.get("stream"):
@@ -494,28 +715,40 @@ def make_handler(svc: Service):
self._sse()
try:
for c in chunks:
self.wfile.write(b"data: " + json.dumps(c, ensure_ascii=False).encode() + b"\n\n")
if c is None:
self.wfile.write(b": keep-alive\n\n") # an SSE comment: clients ignore it
else:
self.wfile.write(b"data: " + json.dumps(c, ensure_ascii=False).encode() + b"\n\n")
self.wfile.flush()
self.wfile.write(b"data: [DONE]\n\n")
except OSError:
cancel.set() # client went away: stop at the next step
cancel.set() # client went away: stop the engine
chunks.close()
def _anthropic(self, req):
messages, tools, kw = anthropic_to_messages(req)
max_new = int(req.get("max_tokens") or 1024)
if max_new <= 0: # a non-positive budget (some clients send -1) means "unset"
max_new = 1024
ids, thinking = svc.prepare(messages, tools, kw, max_new)
_debug_req("anthropic", req, messages, tools, max_new, thinking, len(ids))
cancel = threading.Event()
events = anthropic_events(svc, req, ids, thinking, tools, max_new, cancel)
if not req.get("stream"):
return self._json(200, anthropic_collect(events))
self._sse()
try:
for name, e in events:
self.wfile.write(f"event: {name}\n".encode() + b"data: " +
json.dumps(e, ensure_ascii=False).encode() + b"\n\n")
for item in events:
if item is None:
self.wfile.write(b": keep-alive\n\n")
else:
name, e = item
self.wfile.write(f"event: {name}\n".encode() + b"data: " +
json.dumps(e, ensure_ascii=False).encode() + b"\n\n")
self.wfile.flush()
except OSError:
cancel.set()
events.close()
return Handler
@@ -546,7 +779,7 @@ def main() -> int:
ap.add_argument("--api-key", default=os.environ.get("STRATA_API_KEY", ""),
help="require this key on /v1/* (Authorization: Bearer ... or x-api-key); also $STRATA_API_KEY")
a = ap.parse_args()
cfg = json.loads(Path(a.config).read_text(encoding="utf-8")) if a.config else {}
cfg = json.loads(Path(a.config).read_text(encoding="utf-8-sig")) if a.config else {} # Notepad adds a BOM
try: # before the minutes of loading: is the port free?
Server((a.host, a.port), BaseHTTPRequestHandler).server_close()
except OSError:
+67 -15
View File
@@ -19,7 +19,7 @@ What the first run does (each step is skipped when it is already done):
6. prepares the model for Strata and fetches the MTP draft layer (~5 GB, from the original Qwen checkpoint)
7. writes run-<model>.bat / run-<model>.sh and starts the model
Options: --family qwen|swift, --model Q2_0|IQ2_XS|IQ3_XXS, --context 32768, --vision yes|no|gpu|cpu, --port 8080, --yes (recommended
Options: --family qwen|swift, --model Q2_0|IQ2_XS|IQ3_XXS|IQ3_S, --context 32768, --vision yes|no|gpu|cpu, --port 8080, --yes (recommended
answers, no questions), --setup (install another model / change settings instead of starting), --no-start,
--models-dir DIR, --gguf-dir DIR (use GGUF files you already have), --build (compile instead of the ready-made
engine), --check (only check this PC).
@@ -54,15 +54,19 @@ 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, 1) # the ready-made engine that reads split models (Swift 1.5)
PY_PACKAGES = ["numpy", "jinja2", "regex", "pyyaml", "tqdm", "requests", "cmake", "ninja"]
MIN_ENGINE = (0, 1, 6) # KV streaming (--kv-resident), v0.1.5; its drafter fallback, v0.1.6
PY_PACKAGES = ["numpy", "jinja2", "regex", "pyyaml", "tqdm", "requests", "cmake", "ninja", "pillow"]
MODELS = {
"Q2_0": {"about": "2-bit, the fastest", "download_gb": 66.4, "ram_gb": 48, "arena_gb": 34.0},
"IQ2_XS": {"about": "2-bit i-quant, a little better quality, close in speed", "download_gb": 68.0, "ram_gb": 48,
"arena_gb": 35.5},
"IQ3_XXS": {"about": "3-bit i-quant, the best quality, slower (more CPU work per token)", "download_gb": 75.8,
"IQ3_XXS": {"about": "3-bit i-quant, better quality, slower (more CPU work per token)", "download_gb": 75.8,
"ram_gb": 60, "arena_gb": 42.9},
# the original model only (Swift 1.5 has no IQ3_S): matches the full BF16 model on the published benchmarks
"IQ3_S": {"about": "3.5-bit i-quant, the best quality (matches the full model), the slowest; needs a 64 GB PC "
"with little else running", "download_gb": 83.6, "ram_gb": 62, "arena_gb": 50.3,
"families": ("qwen",)},
}
CONTEXTS = [8192, 32768, 65536, 131072, 262144]
# The model families: the same architecture, weights in the same three GSQ-RCO sizes, different files.
@@ -82,9 +86,9 @@ FAMILIES = {
"license": "Swift Open License 1.0: https://huggingface.co/ukisai/Swift-1.5-Qwen3.8-Flash-Next-GSQ-RCO-GGUF"},
}
MMPROJ = "mmproj-Qwen3.8-Flash-Next-BF16.gguf"
# the image encoder on the GPU: ~0.9 GB of weights + ~0.3 GB of work buffers at 1024 image tokens, kept free of
# expert slots (the engine's default reserve is 700 MiB)
VISION = {"gpu": {"max_tokens": 1024, "reserve_mib": 2100},
# the image encoder on the GPU (~1.2 GB at 1024 image tokens) warms up before the engine starts, so the engine
# sizes its expert slots around it and the default reserve (700 MiB) is enough; engines before 0.1.2 need more
VISION = {"gpu": {"max_tokens": 1024, "reserve_mib": 700},
"cpu": {"max_tokens": 300, "reserve_mib": 700}}
EXE = "strata.exe" if WIN else "strata"
VEXE = "strata-vision.exe" if WIN else "strata-vision"
@@ -456,6 +460,42 @@ def get_prebuilt(url_base, gpu, vision) -> Path | None:
return eng
def update_installed_engine(url_base) -> None:
"""An installed ready-made engine older than MIN_ENGINE is replaced before the model starts, so a plain
START-HERE.bat on an existing install picks up a new release. If that cannot happen (no internet, the model
still running, no ready-made engine for this GPU) the installed engine is kept and starts as before."""
eng = ROOT / "engine"
info = eng / "BUILD.json"
if not info.exists() or not (eng / EXE).exists():
return
meta_text = info.read_text()
meta = json.loads(meta_text)
ver = tuple(int(x) for x in str(meta.get("version", "0")).split(".")[:3] if x.isdigit())
if meta.get("source") == "local" or ver >= MIN_ENGINE:
return
try: # a running engine cannot be replaced (Windows keeps it locked)
for x in (EXE, VEXE):
if (eng / x).exists():
with open(eng / x, "r+b"):
pass
except OSError:
warn(f"engine {meta.get('version')} is in use: close the model window and run this again to update it")
return
gpu = gpu_info()
new = None
if gpu is not None:
try:
new = get_prebuilt(url_base, gpu, "gpu")
except Exception as e: # a failed download must not stop the model from starting
warn(f"updating the engine failed ({e})")
if new is None:
if not info.exists():
info.write_text(meta_text) # get_prebuilt drops it before downloading: put it back
warn(f"could not update the engine: starting the installed {meta.get('version')}")
return
pip_install(CUDA_WHEELS, "NVIDIA CUDA libraries (cuBLAS, CUDA runtime; ~0.4 GB)")
def install_build_tools(gpu, yes):
"""The compiler and the CUDA toolkit, installed for the user (asks once). Returns (nvcc, vcvars)."""
nvcc, cuda_v = find_nvcc()
@@ -572,7 +612,7 @@ def installed_configs():
def start(cfg_path: Path, port: int | None, open_browser=True) -> int:
cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
cfg = 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")
@@ -627,11 +667,13 @@ def main() -> int:
# ---- 0. already installed: just start it
have = installed_configs()
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], None)
say()
for i, c in enumerate(have, 1):
say(f" {i}) {json.loads(c.read_text(encoding='utf-8')).get('model_name', c.stem)}")
say(f" {i}) {json.loads(c.read_text(encoding='utf-8-sig')).get('model_name', c.stem)}")
say(f" {len(have) + 1}) install another model / change settings")
pick = int(ask("Which one?", [str(i) for i in range(1, len(have) + 2)], "1", a.yes))
if pick <= len(have):
@@ -681,13 +723,15 @@ def main() -> int:
if fam.get("license"):
say(f" Its license: {fam['license']}")
say()
names = list(MODELS)
names = [m for m in MODELS if family in MODELS[m].get("families", FAMILIES)]
if a.model and a.model not in names:
fail(f"{fam['title']} has no {a.model} model file", "choose one of: " + ", ".join(names))
for i, m in enumerate(names, 1):
d = MODELS[m]
fit = "" if ram >= d["ram_gb"] else f" <- needs {d['ram_gb']} GB RAM, you have {ram:.0f}"
say(f" {i}) {m:8s} {d['about']}; download {d['download_gb']:.0f} GB, uses ~{d['arena_gb']:.0f} GB of RAM{fit}")
rec = "3" if ram >= 60 else "1"
model = a.model or names[int(ask("Which size?", ["1", "2", "3"], rec, a.yes)) - 1]
rec = str(names.index("IQ3_XXS") + 1) if ram >= 60 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]
if ram < MODELS[model]["ram_gb"] - 4:
fail(f"{model} needs about {MODELS[model]['ram_gb']} GB of RAM; this PC has {ram:.0f} GB",
"choose Q2_0 or IQ2_XS, or add RAM")
@@ -703,8 +747,9 @@ def main() -> int:
for i, c in enumerate(CONTEXTS, 1):
say(f" {i}) {c // 1024}K tokens" + (" (recommended for your GPU)" if c == rec_ctx else ""))
ctx = CONTEXTS[int(ask("Context?", [str(i) for i in range(1, 6)], str(CONTEXTS.index(rec_ctx) + 1), a.yes)) - 1]
if model == "IQ3_XXS" and ram < 90 and ctx > 131072:
warn("IQ3_XXS with a 262K context needs more than 64 GB of RAM (43 GB of experts + the context): using 128K")
if model in ("IQ3_XXS", "IQ3_S") and ram < 90 and ctx > 131072:
warn(f"{model} with a 262K context needs more than 64 GB of RAM ({MODELS[model]['arena_gb']:.0f} GB of experts "
"+ the context): using 128K")
ctx = 131072
ok(f"context: {ctx} tokens")
if a.vision:
@@ -725,7 +770,7 @@ def main() -> int:
# ---- 3. python packages
step(3, "Python packages")
pip_install(PY_PACKAGES, "numpy, jinja2, regex, pyyaml, tqdm, requests, cmake, ninja")
pip_install(PY_PACKAGES, "numpy, jinja2, regex, pyyaml, tqdm, requests, cmake, ninja, pillow")
# ---- 4. the engine
step(4, "the Strata engine")
@@ -818,6 +863,13 @@ def main() -> int:
"--max-context", str(ctx)]
if ctx > 8192:
args += ["--kv", "int8"]
# KV streaming: from 64K up the whole KV cache lives in RAM and only the part the attention reads (32K positions
# per layer) stays in VRAM; the VRAM it frees holds more experts (+6% at 128K, +23% at 262K with Q2_0). It
# costs ~13.7 KB of RAM per context token (1.7 GB at 128K, 3.4 GB at 262K), so only when that fits.
kv_ram_gb = ctx * 13728 / 1e9
if ctx >= 65536 and ram >= MODELS[model]["ram_gb"] + kv_ram_gb + 1:
args += ["--kv-resident", "32768"]
ok(f"KV streaming on: the context's KV cache lives in RAM ({kv_ram_gb:.1f} GB), more experts fit in VRAM")
if vision != "none":
args += ["--vision", "--vram-reserve-mib", str(VISION[vision]["reserve_mib"])]
cfg = {"exe": str(eng / EXE), "args": args, "cwd": str(ROOT), "tokenizer": str(pack / "tokenizer"),
+11
View File
@@ -13,6 +13,7 @@
#include <atomic>
#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <thread>
@@ -263,8 +264,14 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
if (native && d.nact_multi.size() < (size_t) MAXT * kNativeActBytes) d.nact_multi.resize((size_t) MAXT * kNativeActBytes);
if (d.job_of.size() != (size_t) d.n_expert) d.job_of.assign((size_t) d.n_expert, (int16_t) -1);
if (d.jobs_multi.size() < (size_t) (n_tok * k)) d.jobs_multi.resize((size_t) (MAXT * k));
static const bool ptrace = std::getenv("STRATA_POOL_TRACE") != nullptr;
auto pt = [&](const char* what, long long a = -1) {
if (ptrace) { std::fprintf(stderr, "pool trace: layer %lld %s %lld\n", (long long) d.layers, what, a); std::fflush(stderr); }
};
const auto c0 = std::chrono::steady_clock::now();
pt("begin");
d.src->begin_layer(d.layers, ids, n_tok * k);
pt("begun");
if (!d.usage.empty())
for (int64_t i = 0; i < n_tok * k; ++i)
if (ids[i] >= 0 && ids[i] < d.n_expert) d.usage[(size_t) d.layers * (size_t) d.n_expert + (size_t) ids[i]] += 1.0f;
@@ -348,7 +355,9 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
P.counts[1] = entries;
P.counts[2] = fetches;
std::atomic_thread_fence(std::memory_order_seq_cst);
pt("publish", fetches);
if (P.publish) P.publish(P.ctx);
pt("fetch", fetches);
if (P.fetch) P.fetch(P.ctx, dma_src, P.pcie_mode != 0 ? 0 : fetches, (size_t) bb); // the copy engine, beside the CPU's work
} else {
for (int64_t i = 0; i < n; ++i) {
@@ -409,9 +418,11 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
++d.multi_entries;
}
const auto c3 = std::chrono::steady_clock::now();
pt("run", njobs);
if (native) d.pool->run_split_multi_native(lay.fmt[(size_t) d.layers], d.jobs_multi.data(), njobs);
else d.pool->run_split_multi(d.jobs_multi.data(), njobs);
const auto c4 = std::chrono::steady_clock::now();
pt("ran");
auto ms = [](auto a, auto b) { return std::chrono::duration<double, std::milli>(b - a).count(); };
d.ms_plan += ms(c0, c1);
d.ms_actq += ms(c1, c2);
+205 -33
View File
@@ -27,6 +27,8 @@
#include "strata/kernels/qsa_decode_attn.hpp"
#include "strata/kernels/fused_gdn.hpp"
#include "strata/kernels/qsa_select.hpp"
#include <algorithm>
#include <cstdlib>
#include <exception>
#include <cuda_runtime.h>
#include <cmath>
@@ -475,39 +477,207 @@ struct Cursor { uint8_t* p; uint64_t used = 0; template <typename T>
// dispatch too and the two must not be able to drift.
}
// namespace
uint64_t qsa_state_bytes(const ModelGeometry& g, int64_t max_cells, bool with_rope) { const QsaShapes s = qsa_shapes(g); const int64_t pages = (max_cells + s.page_size - 1) / s.page_size; uint64_t n = 0; n += g_kv_int8 ? (uint64_t) pages * s.page_size * strata::kernels::kv_q8_bytes_per_cell(s) + 64 : (uint64_t) pages * s.n_head_kv * s.page_size * s.head_dim * 2 * 2;
// k_pool + v_pool, fp16
n += (uint64_t) pages * 4;
// page_table
n += (uint64_t) (s.idx_block - 1) * s.idx_dim * 4;
// tail
n += (uint64_t) s.idx_dim * 4;
// dead
n += (uint64_t) ((max_cells / s.idx_block) + 2) * s.idx_dim * 4;
// pooled
n += 16;
// block_pos
if (with_rope) n += (uint64_t) max_cells * (s.n_rot / 2) * 4 * 2;
// cos + sin tables
n += strata::kernels::qsa_step_bytes() + 16; // counts and aligned attention status
// the per-token counts
n += (uint64_t) s.n_head * 4;
// pos_dev
return align_up16(n) + 256;}
uint64_t qsa_state_init(const ModelGeometry& g, int64_t max_cells, void* base, QsaState& st, const QsaState* share_rope) { const QsaShapes s = qsa_shapes(g); const int64_t pages = (max_cells + s.page_size - 1) / s.page_size; Cursor c{(uint8_t*) base}; st.kv_int8 = g_kv_int8; if (g_kv_int8) { const uint64_t cells = (uint64_t) pages * s.n_head_kv * s.page_size; st.k_q = c.take<int8_t>(cells * s.head_dim); st.v_q = c.take<int8_t>(cells * s.head_dim); st.k_scale = c.take<uint16_t>(cells * (s.head_dim / strata::kernels::KV_Q8_GROUP)); st.v_scale = c.take<uint16_t>(cells * (s.head_dim / strata::kernels::KV_Q8_GROUP)); } else { st.k_pool = c.take<uint16_t>((uint64_t) pages * s.n_head_kv * s.page_size * s.head_dim); st.v_pool = c.take<uint16_t>((uint64_t) pages * s.n_head_kv * s.page_size * s.head_dim); } st.page_table = c.take<int32_t>((uint64_t) pages); st.n_pages = pages; st.max_cells = max_cells; st.idx_tail = c.take<float>((uint64_t) (s.idx_block - 1) * s.idx_dim); st.idx_dead = c.take<float>((uint64_t) s.idx_dim); st.idx_pooled = c.take<float>((uint64_t) ((max_cells / s.idx_block) + 2) * s.idx_dim); st.idx_block_pos = c.take<int32_t>(1); if (share_rope != nullptr) { st.cos_tab = share_rope->cos_tab; st.sin_tab = share_rope->sin_tab; } else { st.cos_tab = c.take<float>((uint64_t) max_cells * (s.n_rot / 2)); st.sin_tab = c.take<float>((uint64_t) max_cells * (s.n_rot / 2)); } st.step = c.take<int32_t>((uint64_t) strata::kernels::kStepCount); st.attention_status = c.take<int32_t>(1); st.pos_dev = c.take<int32_t>((uint64_t) s.n_head);
// the pinned staging the uploads copy FROM - see the note on `host_step` in the header
if (cudaHostAlloc((void**) &st.host_step, strata::kernels::qsa_step_bytes() + sizeof(int32_t), cudaHostAllocMapped | cudaHostAllocPortable) != cudaSuccess || cudaHostAlloc((void**) &st.host_pos, (size_t) s.n_head * 4, cudaHostAllocMapped | cudaHostAllocPortable) != cudaSuccess) { return 0;
// the caller sees a zero byte count; a half-built state is worse than none
namespace {
// KV streaming (docs/kv-streaming-design.md): 0 keeps every cell in VRAM.
int64_t g_kv_resident = 0;
uint64_t g_kv_host_bytes = 0;
/// How one state holds its K/V: `mode` as in QsaState::kv_mode, `slots` VRAM pages of `pages` logical ones.
struct KvPlan {
int mode = 0;
int64_t pages = 0, slots = 0, pooled_rows = 0;
};
KvPlan kv_plan(const QsaShapes& s, int64_t max_cells, int64_t ring_cells) {
KvPlan p;
p.pages = (max_cells + s.page_size - 1) / s.page_size;
p.slots = p.pages;
p.pooled_rows = max_cells / s.idx_block + 2;
if (g_kv_resident <= 0 || ring_cells < 0) return p; // ring_cells < 0: always fully resident
// A/B only: STRATA_KV_RING_OFF keeps the drafter fully resident, STRATA_KV_MAIN_OFF the main layers
static const bool ring_off = std::getenv("STRATA_KV_RING_OFF") != nullptr;
static const bool main_off = std::getenv("STRATA_KV_MAIN_OFF") != nullptr;
if (ring_cells > 0 ? ring_off : main_off) return p;
if (ring_cells > 0) {
const int64_t r = (ring_cells + s.page_size - 1) / s.page_size;
if (r < p.pages) { p.mode = 2; p.slots = r; p.pooled_rows = 2; } // the drafter has no indexer
} else {
const int64_t r = (std::max(g_kv_resident, qsa_kv_resident_min()) + s.page_size - 1) / s.page_size;
if (r < p.pages) { p.mode = 1; p.slots = r; }
}
return p;
}
st.host_step[strata::kernels::kStepCount] = 0;
// THE ROPE TABLE IS BUILT ON THE HOST IN FLOAT64 and uploaded once, because the reference computes its
// frequencies in float64 and reproducing that on device means double-precision `pow`/`cos` that need not
// agree with the host's libm. A table shorter than the sequence would have the rotation read past it.
if (share_rope == nullptr) { std::vector<float> hc((size_t) max_cells * (s.n_rot / 2)), hs((size_t) max_cells * (s.n_rot / 2)); strata::kernels::build_rope_table((int) s.n_rot, strata::kernels::qsa_freq_base(), (int) max_cells, hc.data(), hs.data()); cudaMemcpy(st.cos_tab, hc.data(), hc.size() * 4, cudaMemcpyHostToDevice); cudaMemcpy(st.sin_tab, hs.data(), hs.size() * 4, cudaMemcpyHostToDevice); }
// the page table starts as the IDENTITY, which is the simplest legal mapping and what a caller that does
// not page at all wants. A real allocator re-points it; nothing outside `kv_append`/`kv_gather` may care.
std::vector<int32_t> tab((size_t) pages); for (int64_t i = 0; i < pages; ++i) tab[(size_t) i] = (int32_t) i; cudaMemcpy(st.page_table, tab.data(), tab.size() * 4, cudaMemcpyHostToDevice); return c.used;}
void qsa_state_zero(const QsaState& st, const ModelGeometry& g, void* stream) { const QsaShapes s = qsa_shapes(g); cudaStream_t cs = (cudaStream_t) stream; if (st.kv_int8) { const size_t cells = (size_t) st.n_pages * s.n_head_kv * s.page_size; cudaMemsetAsync(st.k_q, 0, cells * s.head_dim, cs); cudaMemsetAsync(st.v_q, 0, cells * s.head_dim, cs); cudaMemsetAsync(st.k_scale, 0, cells * (s.head_dim / strata::kernels::KV_Q8_GROUP) * 2, cs); cudaMemsetAsync(st.v_scale, 0, cells * (s.head_dim / strata::kernels::KV_Q8_GROUP) * 2, cs); } else { cudaMemsetAsync(st.k_pool, 0, (size_t) st.n_pages * s.n_head_kv * s.page_size * s.head_dim * 2, cs); cudaMemsetAsync(st.v_pool, 0, (size_t) st.n_pages * s.n_head_kv * s.page_size * s.head_dim * 2, cs); } cudaMemsetAsync(st.idx_tail, 0, (size_t) (s.idx_block - 1) * s.idx_dim * 4, cs); cudaMemsetAsync(st.idx_dead, 0, (size_t) s.idx_dim * 4, cs); cudaMemsetAsync(st.idx_pooled, 0, (size_t) ((st.max_cells / s.idx_block) + 2) * s.idx_dim * 4, cs);}
uint64_t kv_pool_bytes(const QsaShapes& s, int64_t pages) {
return g_kv_int8 ? (uint64_t) pages * s.page_size * strata::kernels::kv_q8_bytes_per_cell(s) + 64
: (uint64_t) pages * s.n_head_kv * s.page_size * s.head_dim * 2 * 2;
}
} // namespace
void qsa_set_kv_resident(int64_t cells) { g_kv_resident = cells > 0 ? cells : 0; }
int64_t qsa_kv_resident() { return g_kv_resident; }
int64_t qsa_kv_resident_min() { return 20480; }
uint64_t qsa_kv_host_bytes() { return g_kv_host_bytes; }
uint64_t qsa_state_bytes(const ModelGeometry& g, int64_t max_cells, bool with_rope, int64_t ring_cells) {
const QsaShapes s = qsa_shapes(g);
const KvPlan p = kv_plan(s, max_cells, ring_cells);
uint64_t n = 0;
n += kv_pool_bytes(s, p.slots) + 4 * 16; // K/V pools (the VRAM slots)
n += (uint64_t) p.pages * 4; // page_table
if (p.mode == 1) n += strata::kernels::kv_stream_map_bytes(p.slots) + 6 * 16; // the residency map
n += (uint64_t) (s.idx_block - 1) * s.idx_dim * 4; // tail
n += (uint64_t) s.idx_dim * 4; // dead
n += (uint64_t) p.pooled_rows * s.idx_dim * 4; // pooled
n += 16; // block_pos
if (with_rope) n += (uint64_t) max_cells * (s.n_rot / 2) * 4 * 2; // cos + sin tables
n += strata::kernels::qsa_step_bytes() + 16; // counts and aligned attention status
n += (uint64_t) s.n_head * 4; // pos_dev
return align_up16(n) + 256;
}
uint64_t qsa_state_init(const ModelGeometry& g, int64_t max_cells, void* base, QsaState& st,
const QsaState* share_rope, int64_t ring_cells) {
const QsaShapes s = qsa_shapes(g);
const KvPlan p = kv_plan(s, max_cells, ring_cells);
const int64_t pages = p.pages;
Cursor c{(uint8_t*) base};
st.kv_int8 = g_kv_int8;
st.kv_mode = p.mode;
st.n_slots = p.slots;
const uint64_t rows = (uint64_t) p.slots * s.n_head_kv * s.page_size; // VRAM rows: the slots
if (g_kv_int8) {
st.k_q = c.take<int8_t>(rows * s.head_dim);
st.v_q = c.take<int8_t>(rows * s.head_dim);
st.k_scale = c.take<uint16_t>(rows * (s.head_dim / strata::kernels::KV_Q8_GROUP));
st.v_scale = c.take<uint16_t>(rows * (s.head_dim / strata::kernels::KV_Q8_GROUP));
} else {
st.k_pool = c.take<uint16_t>(rows * s.head_dim);
st.v_pool = c.take<uint16_t>(rows * s.head_dim);
}
st.page_table = c.take<int32_t>((uint64_t) pages);
st.n_pages = pages;
st.max_cells = max_cells;
st.map = strata::kernels::KvStreamMap{};
st.map.page_table = st.page_table;
st.map.n_blocks = pages;
st.map.n_slots = p.slots;
if (p.mode == 1) {
st.map.slot_block = c.take<int32_t>((uint64_t) p.slots);
st.map.slot_stamp = c.take<int32_t>((uint64_t) p.slots);
st.map.slot_ref = c.take<int32_t>((uint64_t) p.slots);
st.map.miss_block = c.take<int32_t>((uint64_t) p.slots);
st.map.miss_slot = c.take<int32_t>((uint64_t) p.slots);
st.map.ctl = c.take<int32_t>((uint64_t) strata::kernels::kKvCtlInts);
}
st.idx_tail = c.take<float>((uint64_t) (s.idx_block - 1) * s.idx_dim);
st.idx_dead = c.take<float>((uint64_t) s.idx_dim);
st.idx_pooled_rows = p.pooled_rows;
st.idx_pooled = c.take<float>((uint64_t) p.pooled_rows * s.idx_dim);
st.idx_block_pos = c.take<int32_t>(1);
if (share_rope != nullptr) {
st.cos_tab = share_rope->cos_tab;
st.sin_tab = share_rope->sin_tab;
} else {
st.cos_tab = c.take<float>((uint64_t) max_cells * (s.n_rot / 2));
st.sin_tab = c.take<float>((uint64_t) max_cells * (s.n_rot / 2));
}
st.step = c.take<int32_t>((uint64_t) strata::kernels::kStepCount);
st.attention_status = c.take<int32_t>(1);
st.pos_dev = c.take<int32_t>((uint64_t) s.n_head);
// the pinned staging the uploads copy FROM - see the note on `host_step` in the header
if (cudaHostAlloc((void**) &st.host_step, strata::kernels::qsa_step_bytes() + sizeof(int32_t),
cudaHostAllocMapped | cudaHostAllocPortable) != cudaSuccess ||
cudaHostAlloc((void**) &st.host_pos, (size_t) s.n_head * 4, cudaHostAllocMapped | cudaHostAllocPortable) !=
cudaSuccess) {
return 0; // the caller sees a zero byte count; a half-built state is worse than none
}
st.host_step[strata::kernels::kStepCount] = 0;
// KV streaming: the authoritative K/V of every cell, pinned and device-mapped, in the identity layout
st.host = strata::kernels::KvHostPools{};
if (p.mode != 0) {
const uint64_t hrows = (uint64_t) pages * s.n_head_kv * s.page_size;
const uint64_t bytes = (uint64_t) pages * strata::kernels::kv_block_bytes(s, g_kv_int8) + 4 * 256;
uint8_t* h = nullptr;
uint8_t* d = nullptr;
if (cudaHostAlloc((void**) &h, bytes, cudaHostAllocMapped | cudaHostAllocPortable) != cudaSuccess ||
cudaHostGetDevicePointer((void**) &d, h, 0) != cudaSuccess) {
if (p.mode == 1) std::fprintf(stderr, "strata: KV streaming: cannot pin %.2f GiB of RAM for a layer's KV copy "
"(lower the context, or run without --kv-resident)\n", (double) bytes / 1073741824.0);
return 0;
}
g_kv_host_bytes += bytes;
Cursor hc{d};
if (g_kv_int8) {
st.host.k_q = hc.take<int8_t>(hrows * s.head_dim);
st.host.v_q = hc.take<int8_t>(hrows * s.head_dim);
st.host.k_scale = hc.take<uint16_t>(hrows * (s.head_dim / strata::kernels::KV_Q8_GROUP));
st.host.v_scale = hc.take<uint16_t>(hrows * (s.head_dim / strata::kernels::KV_Q8_GROUP));
} else {
st.host.k_pool = hc.take<uint16_t>(hrows * s.head_dim);
st.host.v_pool = hc.take<uint16_t>(hrows * s.head_dim);
}
}
// THE ROPE TABLE IS BUILT ON THE HOST IN FLOAT64 and uploaded once, because the reference computes its
// frequencies in float64 and reproducing that on device means double-precision `pow`/`cos` that need not
// agree with the host's libm. A table shorter than the sequence would have the rotation read past it.
if (share_rope == nullptr) {
std::vector<float> hc((size_t) max_cells * (s.n_rot / 2)), hs((size_t) max_cells * (s.n_rot / 2));
strata::kernels::build_rope_table((int) s.n_rot, strata::kernels::qsa_freq_base(), (int) max_cells,
hc.data(), hs.data());
cudaMemcpy(st.cos_tab, hc.data(), hc.size() * 4, cudaMemcpyHostToDevice);
cudaMemcpy(st.sin_tab, hs.data(), hs.size() * 4, cudaMemcpyHostToDevice);
}
// the page table starts as the IDENTITY, which is the simplest legal mapping and what a caller that does
// not page at all wants; a streamed state starts with nothing resident, a ring at `block % n_slots`.
if (p.mode == 0) {
std::vector<int32_t> tab((size_t) pages);
for (int64_t i = 0; i < pages; ++i) tab[(size_t) i] = (int32_t) i;
cudaMemcpy(st.page_table, tab.data(), tab.size() * 4, cudaMemcpyHostToDevice);
} else if (p.mode == 1) {
strata::kernels::kv_stream_reset(st.map, nullptr);
} else {
strata::kernels::kv_ring_table(st.page_table, pages, p.slots, nullptr);
}
cudaDeviceSynchronize();
return c.used;
}
void qsa_state_zero(const QsaState& st, const ModelGeometry& g, void* stream) {
const QsaShapes s = qsa_shapes(g);
cudaStream_t cs = (cudaStream_t) stream;
const size_t rows = (size_t) st.n_slots * s.n_head_kv * s.page_size;
if (st.kv_int8) {
cudaMemsetAsync(st.k_q, 0, rows * s.head_dim, cs);
cudaMemsetAsync(st.v_q, 0, rows * s.head_dim, cs);
cudaMemsetAsync(st.k_scale, 0, rows * (s.head_dim / strata::kernels::KV_Q8_GROUP) * 2, cs);
cudaMemsetAsync(st.v_scale, 0, rows * (s.head_dim / strata::kernels::KV_Q8_GROUP) * 2, cs);
} else {
cudaMemsetAsync(st.k_pool, 0, rows * s.head_dim * 2, cs);
cudaMemsetAsync(st.v_pool, 0, rows * s.head_dim * 2, cs);
}
// A streamed state starts over with nothing resident. Its host copy is not cleared (GBs over PCIe per new
// conversation): no reader names a cell before this sequence has written it, and a block copied in whole
// carries the unwritten cells past the end, which nothing reads.
if (st.kv_mode == 1) strata::kernels::kv_stream_reset(st.map, stream);
cudaMemsetAsync(st.idx_tail, 0, (size_t) (s.idx_block - 1) * s.idx_dim * 4, cs);
cudaMemsetAsync(st.idx_dead, 0, (size_t) s.idx_dim * 4, cs);
cudaMemsetAsync(st.idx_pooled, 0, (size_t) st.idx_pooled_rows * s.idx_dim * 4, cs);
}
strata::kernels::QsaAttnPools qsa_attn_pools(const QsaState& st) {
strata::kernels::QsaAttnPools pools;
pools.page_table = st.page_table;
if (st.kv_int8) { pools.k_q = st.k_q; pools.v_q = st.v_q; pools.k_scale = st.k_scale; pools.v_scale = st.v_scale; }
else { pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; }
return pools;
}
void qsa_kv_resolve(const QsaState& st, const ModelGeometry& g, const int32_t* ids, const int32_t* steps, int64_t n_q,
int64_t cap, void* stream) {
if (st.kv_mode != 1) return;
strata::kernels::kv_stream_resolve(st.map, qsa_attn_pools(st), st.host, st.kv_int8, ids, steps, n_q, cap,
qsa_shapes(g), stream);
}
uint64_t qsa_buffers_bytes(const ModelGeometry& g, int64_t max_cells) { const QsaShapes s = qsa_shapes(g); const int64_t cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s); uint64_t n = 0; n += (uint64_t) q8k_bytes(g.n_embd); n += (uint64_t) (g.n_embd / 32) * 34;
// block_q8_0
n += (uint64_t) g.n_embd * 2; n += (uint64_t) g.n_head * 2 * g.head_dim * 4; n += (uint64_t) g.n_head * g.head_dim * 4; n += (uint64_t) g.n_head_kv * g.head_dim * 4 * 2; n += (uint64_t) g.idx_key_dim * 4; n += (uint64_t) g.idx_q_heads * g.idx_key_dim * 4; n += (uint64_t) max_cells * 4; n += (uint64_t) cap * 4; n += (uint64_t) cap * g.n_head_kv * g.head_dim * 2 * 2; n += (uint64_t) g.n_head * g.head_dim * 4;
@@ -676,7 +846,7 @@ if (!gemv_quantized(*w_attnk, p_k, f_k, b.x_q8_0, b.x_q8k, b.kcur, g.n_embd, g.n
// EVERY ENTRY POINT FROM HERE ON IS THE CAPTURABLE ONE: the per-token counts come from `st.step` and every
// launch is sized from a capacity in `st`/`b`, so this sequence can be captured and replayed. The
// host-scalar wrappers would be correct here today and silently wrong in a graph.
if (st.kv_int8) kv_append_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, st.step, b.kcur, b.vcur, s, stream); else kv_append_step(st.k_pool, st.v_pool, st.page_table, st.step, b.kcur, b.vcur, s, stream); { const uint64_t nvk = (uint64_t) g.n_head_kv * g.head_dim; const uint64_t base = (uint64_t) 2 * g.n_embd + 2 * g.hc + (uint64_t) g.n_head * g.head_dim + 2 * nvk + 8; if (!st.kv_int8) dump_slot(dump, g, layer, (const float*) st.k_pool, base, nvk / 2, stream); if (!st.kv_int8) dump_slot(dump, g, layer, (const float*) st.v_pool, base + nvk / 2, nvk / 2, stream); dump_slot(dump, g, layer, b.vcur, base + nvk, nvk, stream); dump_slot(dump, g, layer, b.kcur, base + 2 * nvk, nvk, stream); } { const QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos}; if (native_qsa_indexer_enabled()) {
if (st.kv_int8) kv_append_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, st.step, b.kcur, b.vcur, s, stream, &st.host); else kv_append_step(st.k_pool, st.v_pool, st.page_table, st.step, b.kcur, b.vcur, s, stream, &st.host); { const uint64_t nvk = (uint64_t) g.n_head_kv * g.head_dim; const uint64_t base = (uint64_t) 2 * g.n_embd + 2 * g.hc + (uint64_t) g.n_head * g.head_dim + 2 * nvk + 8; if (!st.kv_int8) dump_slot(dump, g, layer, (const float*) st.k_pool, base, nvk / 2, stream); if (!st.kv_int8) dump_slot(dump, g, layer, (const float*) st.v_pool, base + nvk / 2, nvk / 2, stream); dump_slot(dump, g, layer, b.vcur, base + nvk, nvk, stream); dump_slot(dump, g, layer, b.kcur, base + 2 * nvk, nvk, stream); } { const QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos}; if (native_qsa_indexer_enabled()) {
try {
native_qsa_indexer_append(b.idx_raw, st.step + kStepPos, pos_base,
(const float*) w_ikn->data, RMS_EPS, ib, s, st.max_cells, (float) qsa_freq_base(), stream);
@@ -708,6 +878,8 @@ int64_t max_blocks = (st.max_cells / s.idx_block) + 2;
} else {
qsa_index_step(st.idx_pooled, b.q_idx, nullptr, s, st.step, max_blocks, b.cell_scores, stream); topk_512_step(b.cell_scores, s, cap, st.step, b.ids, stream);
}
// KV streaming: every block the selection names is made resident before anything reads it
qsa_kv_resolve(st, g, b.ids, st.step, 1, cap, stream);
if (g_fast_attn && !native_flash_attn_short && dump == nullptr) {
strata::kernels::QsaAttnPools pools;
pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; pools.k_q = st.k_q; pools.v_q = st.v_q;
+29 -9
View File
@@ -157,9 +157,23 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session
// ---- the layer's own K/V (dense attention: no indexer state is read)
const strata::kernels::QsaShapes s = shapes_of(g);
const int64_t max_cells = ss.qsa_states[0].max_cells;
const uint64_t sb = qsa_state_bytes(g, max_cells, false);
// KV streaming: the drafter only reads its last `window` cells, so with streaming on its K/V is a ring of the
// window (plus the cells a round writes ahead of its queries) over a host copy, refilled on a resume. The host copy
// is pinned after the expert arena has pinned what it could: if it does not fit, the K/V stays whole in VRAM.
int64_t ring = (window > 0 && window < max_cells) ? window + 4 * (int64_t) max_t + 64 : 0;
uint64_t sb = qsa_state_bytes(g, max_cells, false, ring);
if (cudaMalloc(&state_arena_, sb) != cudaSuccess) { err = "mtp: the K/V state does not fit"; return false; }
if (qsa_state_init(g, max_cells, state_arena_, st_, &ss.qsa_states[0]) == 0) { err = "mtp: state init failed"; return false; }
if (qsa_state_init(g, max_cells, state_arena_, st_, &ss.qsa_states[0], ring) == 0) {
if (st_.kv_mode == 0) { err = "mtp: state init failed"; return false; }
std::fprintf(stderr, "strata mtp: no pinned RAM left for the draft layer's K/V copy; keeping it in VRAM\n");
cudaGetLastError();
cudaFree(state_arena_);
st_ = QsaState{};
ring = -1; // fully resident
sb = qsa_state_bytes(g, max_cells, false, ring);
if (cudaMalloc(&state_arena_, sb) != cudaSuccess) { err = "mtp: the K/V state does not fit"; return false; }
if (qsa_state_init(g, max_cells, state_arena_, st_, &ss.qsa_states[0], ring) == 0) { err = "mtp: state init failed"; return false; }
}
qsa_state_zero(st_, g, nullptr);
cudaDeviceSynchronize();
vram_ += sb;
@@ -326,10 +340,10 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri
norm_rope(kcur_ + t * NKV * HD, f32("self_attn.k_norm.weight"), (int) NKV, (int) HD, pos + t * NH);
if (st_.kv_int8)
kv_append_q8_step(st_.k_q, st_.v_q, st_.k_scale, st_.v_scale, st_.page_table, step + t * 4,
kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs);
kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs, &st_.host);
else
kv_append_step(st_.k_pool, st_.v_pool, st_.page_table, step + t * 4, kcur_ + t * NKV * HD,
vcur_ + t * NKV * HD, s, cs);
vcur_ + t * NKV * HD, s, cs, &st_.host);
}
if (!full) return true;
native_mmvq(GGML_Q8_0, q8("self_attn.q_proj.weight"), xq_, qfull_, (int) N, (int) (NH * 2 * HD), T, cs);
@@ -342,10 +356,7 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri
}
norm_rope(qc, f32("self_attn.q_norm.weight"), (int) NH, (int) HD, pos + t * NH);
}
QsaAttnPools pools;
pools.page_table = st_.page_table;
if (st_.kv_int8) { pools.k_q = st_.k_q; pools.v_q = st_.v_q; pools.k_scale = st_.k_scale; pools.v_scale = st_.v_scale; }
else { pools.k_pool = st_.k_pool; pools.v_pool = st_.v_pool; }
const QsaAttnPools pools = qsa_attn_pools(st_);
if (window_ > 0) window_ids(const_cast<int32_t*>(step), T, (int) window_, ident_, cap_, cs);
qsa_decode_attn_batch(qcur_, pools, ident_, step, cap_, s, attn_scratch_, attn_, T, cs);
for (int t = 0; t < T; ++t)
@@ -491,6 +502,15 @@ bool MtpDrafter::capture_step(int j, std::string& err) {
return finish_capture(cs_, ok, step_exec_[j], "step", err);
}
void MtpDrafter::kv_restore(int64_t upto) {
if (st_.kv_mode != 2 || upto <= 0) return;
// the ring's blocks below `upto`, from the host copy: a checkpoint resume may have left later cells in them
const strata::kernels::QsaShapes s = shapes_of(*g_);
const int64_t b1 = (upto + s.page_size - 1) / s.page_size, b0 = std::max<int64_t>(0, b1 - st_.n_slots);
strata::kernels::kv_ring_restore(qsa_attn_pools(st_), st_.host, st_.kv_int8, b0, b1, st_.n_slots, s, cs_);
cudaStreamSynchronize(cs_);
}
bool MtpDrafter::prefill(const float* R_rows, const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err) {
const Clock::time_point t0 = Clock::now();
const int64_t HCN = g_->hc * g_->n_embd;
@@ -551,7 +571,7 @@ bool MtpDrafter::draft(int T, const int32_t* tokens, int64_t p, int a, int32_t*
if (probs) probs[0] = pj;
int n = 1;
// the chain continues while the last draft is likely enough to be verified
for (int j = 1; j < max_t_ - 1 && pj >= min_p; ++j) {
for (int j = 1; j < std::min(max_t_ - 1, max_drafts_) && pj >= min_p; ++j) {
if (!capture_step(j, err)) return false;
put(max_t_ + j - 1, p + a + j);
std::atomic_thread_fence(std::memory_order_seq_cst);
+5 -6
View File
@@ -424,10 +424,10 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
const int32_t* step_t = step_ + t * kStepCount;
if (st.kv_int8)
kv_append_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, step_t,
kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs);
kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs, &st.host);
else
kv_append_step(st.k_pool, st.v_pool, st.page_table, step_t, kcur_ + t * NKV * HD,
vcur_ + t * NKV * HD, s, cs);
vcur_ + t * NKV * HD, s, cs, &st.host);
}
const QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos};
for (int t = tb; t < te; ++t)
@@ -454,10 +454,9 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
s, scores_ + (size_t) tb * max_blocks_, cs);
qsa_block_topk(scores_ + (size_t) tb * max_blocks_, step_ + tb * kStepCount, n, max_blocks_, cap_, s,
sel_ + (size_t) tb * cap_, cs);
QsaAttnPools pools;
pools.page_table = st.page_table;
if (st.kv_int8) { pools.k_q = st.k_q; pools.v_q = st.v_q; pools.k_scale = st.k_scale; pools.v_scale = st.v_scale; }
else { pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; }
// KV streaming: the n selections' blocks resident (device-side, inside the graph)
qsa_kv_resolve(st, *g_, sel_ + (size_t) tb * cap_, step_ + tb * kStepCount, n, cap_, cs);
const QsaAttnPools pools = qsa_attn_pools(st);
qsa_decode_attn_batch(qcur_ + tb * NH * HD, pools, sel_ + (size_t) tb * cap_, step_ + tb * kStepCount, cap_,
s, attn_scratch_ + (size_t) tb * attn_scratch_floats_, attn_ + tb * NH * HD, n, cs);
for (int t = tb; t < te; ++t) {
+32 -5
View File
@@ -49,11 +49,12 @@ std::vector<int> physical_cores(bool skip_first) {
#else
cpu_set_t set;
CPU_ZERO(&set);
if (sched_getaffinity(0, sizeof set, &set) == 0)
if (sched_getaffinity(0, sizeof set, &set) == 0) {
for (int i = 0; i < CPU_SETSIZE; ++i)
if (CPU_ISSET(i, &set)) cores.push_back(i);
else
} else {
for (unsigned i = 0; i < std::thread::hardware_concurrency(); ++i) cores.push_back((int) i);
}
#endif
if (skip_first && !cores.empty()) cores.erase(cores.begin());
return cores;
@@ -127,10 +128,23 @@ ExpertPool::ExpertPool(int n_workers, bool pin, bool host_works) : host_works_(h
ExpertPool::~ExpertPool() {
stop_.store(true, std::memory_order_release);
// Bump the epoch so a PARKED worker notices the stop flag rather than sleeping through it.
epoch_.fetch_add(1, std::memory_order_release);
publish();
for (auto& t : threads_) t.join();
}
void ExpertPool::publish() {
// Both sides are seq_cst, and that is the whole lost-wakeup argument: a worker going to sleep does
// `sleepers_++` and then reads `epoch_`, the host does `epoch_++` and then reads `sleepers_`. In one total
// order at least one of them sees the other's write - the worker sees the new epoch and does not sleep, or
// the host sees the sleeper and notifies under the mutex the worker holds until it is inside `wait`.
// On x86 the fetch_add is a locked xadd either way, so this costs the token path nothing.
epoch_.fetch_add(1, std::memory_order_seq_cst);
if (sleepers_.load(std::memory_order_seq_cst) != 0) {
std::lock_guard<std::mutex> lk(sleep_mu_);
sleep_cv_.notify_all();
}
}
void ExpertPool::worker(int id) {
uint32_t seen = 0;
// ARRIVE at the park before the first wait, so `parked_ == n_` is true from construction. Counting only
@@ -147,9 +161,22 @@ void ExpertPool::worker(int id) {
// - a locked read-modify-write, five workers against one cache line - so the workers spent their wait
// invalidating each other's caches and the very line the host writes to publish work. The counter was
// diagnostic and nothing branched on it. See the note on the atomics in pool.hpp.
//
// After `kSpinBeforeSleep` with no work the worker sleeps instead (issue #4). The clock is read once
// every 1024 pauses, so the spin itself is unchanged.
const auto parked_at = std::chrono::steady_clock::now();
uint32_t spins = 0;
while (epoch_.load(std::memory_order_acquire) == seen) {
if (stop_.load(std::memory_order_relaxed)) return;
_mm_pause();
if ((++spins & 1023u) != 0) continue;
if (std::chrono::steady_clock::now() - parked_at < kSpinBeforeSleep) continue;
std::unique_lock<std::mutex> lk(sleep_mu_);
sleepers_.fetch_add(1, std::memory_order_seq_cst);
sleep_cv_.wait(lk, [&] {
return epoch_.load(std::memory_order_seq_cst) != seen || stop_.load(std::memory_order_relaxed);
});
sleepers_.fetch_sub(1, std::memory_order_relaxed);
}
if (stop_.load(std::memory_order_acquire)) return;
seen = epoch_.load(std::memory_order_relaxed);
@@ -245,7 +272,7 @@ void ExpertPool::run_phase(int mode, int n_tasks) {
njobs_ = n_tasks;
head_.store(0, std::memory_order_relaxed);
done_.store(0, std::memory_order_relaxed);
epoch_.fetch_add(1, std::memory_order_release);
publish();
if (host_works_) drain(-1, host_scratch_);
while (done_.load(std::memory_order_acquire) != (uint32_t) n_tasks) _mm_pause();
while (parked_.load(std::memory_order_acquire) != (uint32_t) n_) _mm_pause();
@@ -356,7 +383,7 @@ void ExpertPool::run(ExpertJob* jobs, int n) {
mode_ = 0;
head_.store(0, std::memory_order_relaxed);
done_.store(0, std::memory_order_relaxed);
epoch_.fetch_add(1, std::memory_order_release); // release: jobs_/njobs_ are visible before the bump
publish(); // a release (seq_cst): jobs_/njobs_ are visible before the bump
// ---- **THE HOST DRAINS TOO (R2.2), INSTEAD OF SPINNING ON `done_`.**
//
+27
View File
@@ -268,6 +268,28 @@ __device__ __forceinline__ float vec_dot_iq4_nl_q8_1(const void* __restrict__ vb
return d * sumi;
}
// IQ4_XS: 256 values as 8 sub-blocks of 32 (6-bit scale each); one call covers one sub-block (iqs = 4 * sub-block),
// and `bq8_1` is the super-block's first q8_1 block, so the call's activation is bq8_1[iqs / 4]. The GSQ-RCO IQ3_S
// file keeps one layer's routed gate/up experts in this format.
__device__ __forceinline__ float vec_dot_iq4_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1,
const int& kbx, const int& iqs) {
const block_iq4_xs* bq4 = (const block_iq4_xs*) vbq + kbx;
int sumi = 0;
#pragma unroll
for (int j = 0; j < 4; ++j) {
const int aux_q4 = get_int_b4(bq4->qs, iqs + j);
const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl);
const int u0 = get_int_b4(bq8_1[iqs / 4].qs, j + 0);
const int u1 = get_int_b4(bq8_1[iqs / 4].qs, j + 4);
sumi = ggml_cuda_dp4a(v.x, u0, sumi);
sumi = ggml_cuda_dp4a(v.y, u1, sumi);
}
const int ls = ((bq4->scales_l[iqs / 8] >> (iqs & 0x04)) & 0x0F) | (((bq4->scales_h >> (iqs / 2)) & 0x03) << 4);
sumi *= ls - 32;
const float d = __half2float(bq4->d) * __low2float(bq8_1[iqs / 4].ds);
return d * sumi;
}
// ---------------------------------------------------------------- the formats
// qk = values per block, ipb = dot calls per block (qi / vdr), step = the iqs stride between calls.
template<int TY> struct Fmt;
@@ -281,6 +303,8 @@ template<> struct Fmt<20> { static constexpr int qk = 32, ipb = 2, step = 2;
__device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq4_nl_q8_1(v, y, kbx, iqs); } };
template<> struct Fmt<21> { static constexpr int qk = 256, ipb = 8, step = 2;
__device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq3_s_q8_1(v, y, kbx, iqs); } };
template<> struct Fmt<23> { static constexpr int qk = 256, ipb = 8, step = 4;
__device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq4_xs_q8_1(v, y, kbx, iqs); } };
template<> struct Fmt<22> { static constexpr int qk = 256, ipb = 8, step = 2;
__device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq2_s_q8_1(v, y, kbx, iqs); } };
template<> struct Fmt<29> { static constexpr int qk = 256, ipb = 8, step = 1;
@@ -618,6 +642,7 @@ void iq_mmvq(int t, const void* w, const void* x_q8_1, float* y, int n_in, int n
case 20: mmvq_kernel<20><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols); break;
case 21: mmvq_kernel<21><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols); break;
case 22: mmvq_kernel<22><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols); break;
case 23: mmvq_kernel<23><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols); break;
case 29: mmvq_kernel<29><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols); break;
case 42: mmvq_kernel<42><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols); break;
default: std::fprintf(stderr, "iq_mmvq: type %d is not supported\n", t); std::exit(1);
@@ -700,6 +725,7 @@ void native_expert_grouped(const NativeExpertLayout& L, const unsigned long long
case 18: native_gu_kernel<18><<<ggu, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break;
case 21: native_gu_kernel<21><<<ggu, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break;
case 22: native_gu_kernel<22><<<ggu, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break;
case 23: native_gu_kernel<23><<<ggu, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break;
case 29: native_gu_kernel<29><<<ggu, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break;
case 42: native_gu_kernel<42><<<ggu, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break;
default: std::fprintf(stderr, "native_expert_grouped: gate/up type %d\n", L.gu_type); std::exit(1);
@@ -711,6 +737,7 @@ void native_expert_grouped(const NativeExpertLayout& L, const unsigned long long
const dim3 gd((unsigned) ((L.n_embd + 7) / 8), (unsigned) cap_groups);
switch (L.d_type) {
case 20: native_down_kernel<20><<<gd, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); break;
case 23: native_down_kernel<23><<<gd, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); break;
case 42: native_down_kernel<42><<<gd, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); break;
default: std::fprintf(stderr, "native_expert_grouped: down type %d\n", L.d_type); std::exit(1);
}
+15 -6
View File
@@ -31,7 +31,7 @@ __global__ void kv_append_q8_kernel(int8_t* __restrict__ k_q, int8_t* __restrict
uint16_t* __restrict__ k_scale, uint16_t* __restrict__ v_scale,
const int32_t* __restrict__ table, const int32_t* __restrict__ step,
const float* __restrict__ kcur, const float* __restrict__ vcur, int kv_heads,
int head_dim, int page_size) {
int head_dim, int page_size, KvHostPools host) {
const long long pos = (long long) __ldg(step + kStepPos);
const int h = blockIdx.x, g = blockIdx.y, t = threadIdx.x;
const bool is_v = blockIdx.z == 1;
@@ -51,10 +51,18 @@ __global__ void kv_append_q8_kernel(int8_t* __restrict__ k_q, int8_t* __restrict
q = __float2int_rn(x / sf);
q = q < -127 ? -127 : (q > 127 ? 127 : q);
}
// KV streaming: the host copy (identity layout) always, the VRAM page only if the block is resident
const long long page = (long long) table[pos / page_size];
const long long row = (page * kv_heads + h) * page_size + (pos % page_size);
(is_v ? v_q : k_q)[row * head_dim + g * KV_Q8_GROUP + t] = (int8_t) q;
if (t == 0) (is_v ? v_scale : k_scale)[row * groups + g] = sbits;
if (page >= 0) {
const long long row = (page * kv_heads + h) * page_size + (pos % page_size);
(is_v ? v_q : k_q)[row * head_dim + g * KV_Q8_GROUP + t] = (int8_t) q;
if (t == 0) (is_v ? v_scale : k_scale)[row * groups + g] = sbits;
}
if (host.k_q != nullptr) {
const long long row = ((pos / page_size) * kv_heads + h) * page_size + (pos % page_size);
(is_v ? host.v_q : host.k_q)[row * head_dim + g * KV_Q8_GROUP + t] = (int8_t) q;
if (t == 0) (is_v ? host.v_scale : host.k_scale)[row * groups + g] = sbits;
}
}
// One thread = 4 consecutive values of one cell and head (as the FP16 gather does with uint2).
@@ -93,12 +101,13 @@ __global__ void kv_gather_q8_kernel(const int8_t* __restrict__ k_q, const int8_t
} // namespace
void kv_append_q8_step(int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, const int32_t* page_table,
const int32_t* step, const float* kcur, const float* vcur, const QsaShapes& s, void* stream) {
const int32_t* step, const float* kcur, const float* vcur, const QsaShapes& s, void* stream,
const KvHostPools* host) {
validate(s, "kv_append_q8");
const dim3 grid((unsigned) s.n_head_kv, (unsigned) (s.head_dim / KV_Q8_GROUP), 2);
kv_append_q8_kernel<<<grid, KV_Q8_GROUP, 0, (cudaStream_t) stream>>>(
k_q, v_q, k_scale, v_scale, page_table, step, kcur, vcur, (int) s.n_head_kv, (int) s.head_dim,
(int) s.page_size);
(int) s.page_size, host ? *host : KvHostPools{});
check("kv_append_q8 launch");
}
+251
View File
@@ -0,0 +1,251 @@
// src/kernels/cuda/kv_stream.cu - see include/strata/kernels/kv_stream.hpp.
#include "strata/kernels/kv_stream.hpp"
#include "strata/kernels/kv_q8.hpp"
#include <cuda_runtime.h>
#include <algorithm>
#include <cstdio>
#include <cstdlib>
namespace strata::kernels {
namespace {
void check(const char* what) {
const cudaError_t e = cudaGetLastError();
if (e != cudaSuccess) {
std::fprintf(stderr, "kv_stream: %s: %s\n", what, cudaGetErrorString(e));
std::exit(1);
}
}
constexpr int RT = 1024; // the resolve block
// The per-block byte runs of the (up to four) pool arrays: block b of array i is bytes [b * len, (b + 1) * len).
struct Runs {
const uint8_t* src[4];
uint8_t* dst[4];
int len[4];
int n;
};
Runs runs_of(const QsaAttnPools& slots, const KvHostPools& host, bool int8, const QsaShapes& s) {
const int rows = (int) (s.n_head_kv * s.page_size);
Runs r{};
if (int8) {
const int codes = rows * (int) s.head_dim, scales = rows * (int) (s.head_dim / KV_Q8_GROUP) * 2;
r.src[0] = (const uint8_t*) host.k_q; r.dst[0] = (uint8_t*) slots.k_q; r.len[0] = codes;
r.src[1] = (const uint8_t*) host.v_q; r.dst[1] = (uint8_t*) slots.v_q; r.len[1] = codes;
r.src[2] = (const uint8_t*) host.k_scale; r.dst[2] = (uint8_t*) slots.k_scale; r.len[2] = scales;
r.src[3] = (const uint8_t*) host.v_scale; r.dst[3] = (uint8_t*) slots.v_scale; r.len[3] = scales;
r.n = 4;
} else {
const int bytes = rows * (int) s.head_dim * 2;
r.src[0] = (const uint8_t*) host.k_pool; r.dst[0] = (uint8_t*) slots.k_pool; r.len[0] = bytes;
r.src[1] = (const uint8_t*) host.v_pool; r.dst[1] = (uint8_t*) slots.v_pool; r.len[1] = bytes;
r.n = 2;
}
return r;
}
// Block-wide exclusive prefix sum of one int per thread (RT threads); `total` is the sum over the block.
__device__ int block_scan(int v, int* warp_sums, int& total) {
const int lane = threadIdx.x & 31, w = threadIdx.x >> 5;
int x = v;
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) warp_sums[w] = x;
__syncthreads();
if (w == 0) {
int t = warp_sums[lane];
for (int o = 1; o < 32; o <<= 1) {
const int y = __shfl_up_sync(0xffffffffu, t, o);
if (lane >= o) t += y;
}
warp_sums[lane] = t;
}
__syncthreads();
total = warp_sums[31];
const int excl = x - v + (w > 0 ? warp_sums[w - 1] : 0);
__syncthreads();
return excl;
}
__global__ void __launch_bounds__(RT) resolve_kernel(KvStreamMap m, const int32_t* __restrict__ ids,
const int32_t* __restrict__ steps, int n_q, int cap, int page_size) {
__shared__ int s_nmiss, s_lookups, s_cut;
__shared__ int warp_sums[32];
const int epoch = m.ctl[0] + 1;
if (threadIdx.x == 0) { s_nmiss = 0; s_lookups = 0; }
__syncthreads();
// 1. hits take this epoch and their reference bit; a missing block is claimed exactly once (-1 -> -2)
int lookups = 0;
for (int q = 0; q < n_q; ++q) {
const int width = steps[q * kStepCount + kStepWidth];
const int32_t* qi = ids + (long long) q * cap;
for (int i = threadIdx.x; i < width; i += RT) {
const int b = qi[i] / page_size;
if (i > 0 && qi[i - 1] / page_size == b) continue; // ids are ascending: one lookup per block
++lookups;
const int sl = m.page_table[b];
if (sl >= 0) {
m.slot_stamp[sl] = epoch;
m.slot_ref[sl] = 1;
} else if (sl == -1 && atomicCAS(&m.page_table[b], -1, -2) == -1) {
m.miss_block[atomicAdd(&s_nmiss, 1)] = b;
}
}
}
atomicAdd(&s_lookups, lookups);
__syncthreads();
// 2. one victim per miss: a clock sweep from the hand. A slot this call uses (stamp == epoch) is never taken;
// a referenced one loses its bit as the hand passes it and is taken on the next pass.
const int need = s_nmiss, n = (int) m.n_slots;
int hand = m.ctl[1], got = 0;
for (int scanned = 0; got < need && scanned < 3 * n; scanned += RT) {
const int j = (int) (((long long) hand + threadIdx.x) % n);
const bool mine = m.slot_stamp[j] == epoch;
const bool cand = !mine && (m.slot_block[j] < 0 || m.slot_ref[j] == 0);
int total = 0;
const int rank = block_scan(cand ? 1 : 0, warp_sums, total);
const int want = need - got;
if (threadIdx.x == 0) s_cut = RT;
__syncthreads();
if (cand && rank == want - 1) s_cut = threadIdx.x + 1; // the hand stops just past the last slot taken
__syncthreads();
const int cut = s_cut;
if (cand && rank < want) {
m.miss_slot[got + rank] = j;
m.slot_stamp[j] = epoch; // taken: a sweep that wraps around must not take it twice
} else if (threadIdx.x < cut && !mine) {
m.slot_ref[j] = 0;
}
got += total < want ? total : want;
hand = (int) (((long long) hand + cut) % n);
__syncthreads();
}
// 3. re-point the table; the copy kernel fills the slots
const int placed = got < need ? got : need;
for (int k = threadIdx.x; k < need; k += RT) {
const int b = m.miss_block[k];
if (k >= placed) { m.page_table[b] = -1; continue; } // overflow: never happens with a legal n_slots
const int sl = m.miss_slot[k];
const int old = m.slot_block[sl];
if (old >= 0) m.page_table[old] = -1;
m.slot_block[sl] = b;
m.slot_stamp[sl] = epoch;
m.slot_ref[sl] = 1;
m.page_table[b] = sl;
}
if (threadIdx.x == 0) {
m.ctl[0] = epoch;
m.ctl[1] = hand;
m.ctl[2] = placed;
if (placed < need) m.ctl[3] = 1;
unsigned long long* c = reinterpret_cast<unsigned long long*>(m.ctl + 4);
c[0] += (unsigned long long) placed;
c[1] += (unsigned long long) s_lookups;
c[2] += 1ull;
}
}
// One block per missed block (grid-stride): copy its runs from the host copy into its slot, 16 B per thread.
__global__ void copy_kernel(KvStreamMap m, Runs r) {
const int need = m.ctl[2];
for (int k = blockIdx.x; k < need; k += gridDim.x) {
const long long b = m.miss_block[k], sl = m.miss_slot[k];
for (int a = 0; a < r.n; ++a) {
const uint4* src = reinterpret_cast<const uint4*>(r.src[a] + b * r.len[a]);
uint4* dst = reinterpret_cast<uint4*>(r.dst[a] + sl * r.len[a]);
for (int i = threadIdx.x; i < r.len[a] / 16; i += blockDim.x) dst[i] = src[i];
}
}
}
__global__ void reset_kernel(KvStreamMap m) {
const long long i0 = (long long) blockIdx.x * blockDim.x + threadIdx.x, st = (long long) gridDim.x * blockDim.x;
for (long long i = i0; i < m.n_blocks; i += st) m.page_table[i] = -1;
for (long long i = i0; i < m.n_slots; i += st) {
m.slot_block[i] = -1;
m.slot_stamp[i] = -1;
m.slot_ref[i] = 0;
}
if (i0 < kKvCtlInts) m.ctl[i0] = 0;
}
__global__ void ring_kernel(int32_t* table, long long n_blocks, long long n_slots) {
for (long long i = (long long) blockIdx.x * blockDim.x + threadIdx.x; i < n_blocks;
i += (long long) gridDim.x * blockDim.x)
table[i] = (int32_t) (i % n_slots);
}
} // namespace
uint64_t kv_block_bytes(const QsaShapes& s, bool int8) {
const uint64_t rows = (uint64_t) (s.n_head_kv * s.page_size);
return int8 ? rows * (uint64_t) s.head_dim * 2 + rows * (uint64_t) (s.head_dim / KV_Q8_GROUP) * 2 * 2
: rows * (uint64_t) s.head_dim * 2 * 2;
}
void kv_stream_reset(const KvStreamMap& m, void* stream) {
reset_kernel<<<128, 256, 0, (cudaStream_t) stream>>>(m);
check("reset");
}
void kv_stream_resolve(const KvStreamMap& m, const QsaAttnPools& slots, const KvHostPools& host, bool int8,
const int32_t* ids, const int32_t* steps, int64_t n_q, int64_t cap, const QsaShapes& s,
void* stream) {
if (n_q <= 0) return;
if (s.n_head_kv * s.page_size * (s.head_dim / KV_Q8_GROUP) * 2 % 16 != 0) {
std::fprintf(stderr, "kv_stream: a block's scale run must be a multiple of 16 bytes\n");
std::exit(1);
}
resolve_kernel<<<1, RT, 0, (cudaStream_t) stream>>>(m, ids, steps, (int) n_q, (int) cap, (int) s.page_size);
check("resolve");
copy_kernel<<<96, 128, 0, (cudaStream_t) stream>>>(m, runs_of(slots, host, int8, s));
check("copy");
}
void kv_ring_table(int32_t* page_table, int64_t n_blocks, int64_t n_slots, void* stream) {
ring_kernel<<<64, 256, 0, (cudaStream_t) stream>>>(page_table, n_blocks, n_slots);
check("ring table");
}
void kv_ring_restore(const QsaAttnPools& slots, const KvHostPools& host, bool int8, int64_t b0, int64_t b1,
int64_t n_slots, const QsaShapes& s, void* stream) {
const Runs r = runs_of(slots, host, int8, s);
for (int64_t b = b0; b < b1;) {
const int64_t sl = b % n_slots, run = std::min<int64_t>(b1 - b, n_slots - sl); // up to the ring's end
for (int a = 0; a < r.n; ++a)
if (cudaMemcpyAsync(r.dst[a] + sl * r.len[a], r.src[a] + b * r.len[a], (size_t) (run * r.len[a]),
cudaMemcpyDefault, (cudaStream_t) stream) != cudaSuccess)
check("ring restore");
b += run;
}
}
void kv_stage_from_host(const QsaAttnPools& stage, const KvHostPools& host, bool int8, int64_t n_blocks,
const QsaShapes& s, void* stream) {
if (n_blocks <= 0) return;
const Runs r = runs_of(stage, host, int8, s);
for (int a = 0; a < r.n; ++a)
if (cudaMemcpyAsync(r.dst[a], r.src[a], (size_t) (n_blocks * r.len[a]), cudaMemcpyDefault,
(cudaStream_t) stream) != cudaSuccess)
check("stage");
}
KvStreamCounters kv_stream_counters(const KvStreamMap& m) {
int32_t c[kKvCtlInts] = {};
KvStreamCounters r;
if (m.ctl == nullptr || cudaMemcpy(c, m.ctl, sizeof(c), cudaMemcpyDeviceToHost) != cudaSuccess) return r;
const unsigned long long* u = reinterpret_cast<const unsigned long long*>(c + 4);
r.misses = u[0];
r.lookups = u[1];
r.calls = u[2];
r.overflow = c[3] != 0;
return r;
}
} // namespace strata::kernels
+17 -7
View File
@@ -43,6 +43,7 @@
#include "strata/kernels/qsa.hpp"
#include "strata/kernels/f16_bits.hpp"
#include "strata/kernels/kv_stream.hpp"
#include "strata/kernels/rope.hpp"
#include "strata/kernels/mrope.hpp"
@@ -120,16 +121,24 @@ __device__ __forceinline__ float h2f(uint16_t bits) { return f32_from_f16(bits);
__global__ void kv_append_kernel(uint16_t* __restrict__ k_pool, uint16_t* __restrict__ v_pool,
const int32_t* __restrict__ table, const int32_t* __restrict__ step,
const float* __restrict__ kcur, const float* __restrict__ vcur, int kv_heads,
int head_dim, int page_size) {
int head_dim, int page_size, KvHostPools host) {
const long long pos = (long long) __ldg(step + kStepPos);
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= kv_heads * head_dim) return;
const int h = i / head_dim, d = i - h * head_dim;
// `[page][kv_head][page_size][head_dim]`: one head's consecutive cells are contiguous inside a page
// `[page][kv_head][page_size][head_dim]`: one head's consecutive cells are contiguous inside a page.
// KV streaming: the host copy (identity layout) always, the VRAM page only if the block is resident.
const long long page = (long long) table[pos / page_size];
const long long row = (page * kv_heads + h) * page_size + (pos % page_size);
k_pool[row * head_dim + d] = f16_from_f32(kcur[i]);
v_pool[row * head_dim + d] = f16_from_f32(vcur[i]);
if (page >= 0) {
const long long row = (page * kv_heads + h) * page_size + (pos % page_size);
k_pool[row * head_dim + d] = f16_from_f32(kcur[i]);
v_pool[row * head_dim + d] = f16_from_f32(vcur[i]);
}
if (host.k_pool != nullptr) {
const long long row = ((pos / page_size) * kv_heads + h) * page_size + (pos % page_size);
host.k_pool[row * head_dim + d] = f16_from_f32(kcur[i]);
host.v_pool[row * head_dim + d] = f16_from_f32(vcur[i]);
}
}
// ================= 2. indexer_key_append =================
@@ -575,13 +584,14 @@ const int32_t* step_upload_width(int64_t width, const QsaShapes& s) {
// buffer and forward this token's counts as the capacities.
void kv_append_step(uint16_t* k_pool, uint16_t* v_pool, const int32_t* page_table, const int32_t* step,
const float* kcur, const float* vcur, const QsaShapes& s, void* stream) {
const float* kcur, const float* vcur, const QsaShapes& s, void* stream, const KvHostPools* host) {
validate(s, "kv_append");
if (step == nullptr) fail("kv_append: step is null");
// The grid is the head x dim count, which is a CONSTANT - `kv_append` writes one cell.
const long long n = s.n_head_kv * s.head_dim;
kv_append_kernel<<<grid_for(n, THREADS), THREADS, 0, (cudaStream_t) stream>>>(
k_pool, v_pool, page_table, step, kcur, vcur, (int) s.n_head_kv, (int) s.head_dim, (int) s.page_size);
k_pool, v_pool, page_table, step, kcur, vcur, (int) s.n_head_kv, (int) s.head_dim, (int) s.page_size,
host ? *host : KvHostPools{});
check_launch("kv_append");
if (stream == nullptr) check_sync("kv_append");
}
+14 -7
View File
@@ -244,12 +244,14 @@ __global__ void ident_hits_kernel(const int32_t* __restrict__ ids, int n, int32_
if (i == 0) *count = n;
}
__global__ void gather_rows_kernel(const uint4* __restrict__ src, long long row16, const int32_t* __restrict__ ids,
long long n, uint4* __restrict__ dst) {
const long long total = n * row16;
// E = the widest element the row size divides into (16, 4 or 1 bytes): a Q6_K head row of 2560 values is 2100 bytes
template<typename E>
__global__ void gather_rows_kernel(const E* __restrict__ src, long long row_e, const int32_t* __restrict__ ids,
long long n, E* __restrict__ dst) {
const long long total = n * row_e;
for (long long i = (long long) blockIdx.x * blockDim.x + threadIdx.x; i < total; i += (long long) gridDim.x * blockDim.x) {
const long long r = i / row16, o = i - r * row16;
dst[i] = src[(long long) ids[r] * row16 + o];
const long long r = i / row_e, o = i - r * row_e;
dst[i] = src[(long long) ids[r] * row_e + o];
}
}
@@ -334,8 +336,13 @@ void mtp_select(const float* R_src, int64_t R_stride, const int32_t* ids, const
}
void gather_rows(const uint8_t* src, int64_t row_bytes, const int32_t* ids, int64_t n, uint8_t* dst, void* stream) {
if (row_bytes % 16 != 0) { std::fprintf(stderr, "gather_rows: row size must be a multiple of 16\n"); std::exit(1); }
gather_rows_kernel<<<48 * 8, 256, 0, (cudaStream_t) stream>>>((const uint4*) src, row_bytes / 16, ids, n, (uint4*) dst);
cudaStream_t s = (cudaStream_t) stream;
if (row_bytes % 16 == 0)
gather_rows_kernel<<<48 * 8, 256, 0, s>>>((const uint4*) src, row_bytes / 16, ids, n, (uint4*) dst);
else if (row_bytes % 4 == 0)
gather_rows_kernel<<<48 * 8, 256, 0, s>>>((const uint32_t*) src, row_bytes / 4, ids, n, (uint32_t*) dst);
else
gather_rows_kernel<<<48 * 8, 256, 0, s>>>(src, row_bytes, ids, n, dst);
check("gather_rows");
}
+227
View File
@@ -0,0 +1,227 @@
// src/kernels/kv_stream_parity.cpp - KV streaming (kv_stream.hpp) against a fully resident pool (GPU, no model).
//
// Decodes a synthetic sequence twice: into a fully resident pool with the identity page table, and into a streamed
// layer (host copy + a small VRAM slot pool that must evict constantly). Every few cells a batch of 1-8 queries
// with selections shaped like the indexer's (whole 4-cell blocks, a partial threshold block, the tail block, half of
// them recent) is attended in both; `kv_stream_resolve` runs before the streamed attention. Checks:
// 1. the attention outputs are BITWISE equal (the streamed readers see exactly the resident values);
// 2. the residency map is consistent after every call (slot_block and page_table invert each other);
// 3. no call overflowed, and the hit/miss counters add up;
// 4. a ring (the MTP drafter's layout) restored from the host copy reads the same values as the resident pool.
// INT8 and FP16 pools both.
#include "strata/kernels/kv_q8.hpp"
#include "strata/kernels/kv_stream.hpp"
#include "strata/kernels/qsa.hpp"
#include "strata/kernels/qsa_decode_attn.hpp"
#include <cuda_runtime.h>
#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <random>
#include <set>
#include <vector>
namespace k = strata::kernels;
namespace {
int g_fail = 0;
void ck(cudaError_t e, const char* w) {
if (e != cudaSuccess) { std::fprintf(stderr, "%s: %s\n", w, cudaGetErrorString(e)); std::exit(2); }
}
template <typename T> T* dalloc(size_t n) {
T* p = nullptr;
ck(cudaMalloc(&p, n * sizeof(T) + 64), "malloc");
ck(cudaMemset(p, 0, n * sizeof(T) + 64), "memset");
return p;
}
template <typename T> T* halloc(size_t n) { // pinned, mapped; returns the device pointer
void* h = nullptr;
void* d = nullptr;
ck(cudaHostAlloc(&h, n * sizeof(T) + 64, cudaHostAllocMapped), "hostalloc");
std::memset(h, 0, n * sizeof(T) + 64);
ck(cudaHostGetDevicePointer(&d, h, 0), "devptr");
return (T*) d;
}
struct Pools { // one K/V pool set of `pages` pages
k::KvHostPools p; // reused as a plain pointer bundle
void alloc(int64_t pages, const k::QsaShapes& s, bool int8, bool host) {
const size_t rows = (size_t) pages * s.n_head_kv * s.page_size;
if (int8) {
p.k_q = host ? halloc<int8_t>(rows * s.head_dim) : dalloc<int8_t>(rows * s.head_dim);
p.v_q = host ? halloc<int8_t>(rows * s.head_dim) : dalloc<int8_t>(rows * s.head_dim);
p.k_scale = host ? halloc<uint16_t>(rows * 4) : dalloc<uint16_t>(rows * 4);
p.v_scale = host ? halloc<uint16_t>(rows * 4) : dalloc<uint16_t>(rows * 4);
} else {
p.k_pool = host ? halloc<uint16_t>(rows * s.head_dim) : dalloc<uint16_t>(rows * s.head_dim);
p.v_pool = host ? halloc<uint16_t>(rows * s.head_dim) : dalloc<uint16_t>(rows * s.head_dim);
}
}
k::QsaAttnPools attn(const int32_t* table) const {
k::QsaAttnPools a;
a.k_pool = p.k_pool; a.v_pool = p.v_pool; a.k_q = p.k_q; a.v_q = p.v_q; a.k_scale = p.k_scale;
a.v_scale = p.v_scale; a.page_table = table;
return a;
}
};
void append(const Pools& pl, const int32_t* table, const int32_t* step, const float* kc, const float* vc,
const k::QsaShapes& s, bool int8, const k::KvHostPools* host) {
if (int8) k::kv_append_q8_step(pl.p.k_q, pl.p.v_q, pl.p.k_scale, pl.p.v_scale, table, step, kc, vc, s, nullptr, host);
else k::kv_append_step(pl.p.k_pool, pl.p.v_pool, table, step, kc, vc, s, nullptr, host);
}
// A selection like qsa_block_topk's: ascending cells, whole blocks, the tail block's cells, and one partial block.
std::vector<int32_t> selection(int64_t n_kv, int64_t width, std::mt19937& rng) {
std::vector<int32_t> ids;
if (n_kv <= width) {
for (int64_t c = 0; c < n_kv; ++c) ids.push_back((int32_t) c);
return ids;
}
const int64_t n_bid = n_kv / 4, tail = n_kv - n_bid * 4;
std::set<int64_t> blocks;
const int64_t want_full = (width - tail) / 4, part = (width - tail) % 4;
std::uniform_int_distribution<int64_t> any(0, n_bid - 1), recent(std::max<int64_t>(0, n_bid - 2048), n_bid - 1);
while ((int64_t) blocks.size() < want_full + (part ? 1 : 0)) blocks.insert(rng() % 2 ? recent(rng) : any(rng));
int64_t partial_block = part ? *std::next(blocks.begin(), (long) (rng() % blocks.size())) : -1;
for (int64_t b : blocks) {
const int64_t take = b == partial_block ? part : 4;
for (int64_t i = 0; i < take; ++i) ids.push_back((int32_t) (b * 4 + i));
}
for (int64_t c = n_bid * 4; c < n_kv; ++c) ids.push_back((int32_t) c);
std::sort(ids.begin(), ids.end());
return ids;
}
bool run(bool int8) {
k::QsaShapes s = k::qsa_real_shapes();
const int64_t N = 40000, n_blocks = (N + 3) / 4, n_slots = 8 * 516 + 700; // must evict: slots < blocks
const int64_t cap = k::qsa_selection_width(k::kTopkMaxCells, s), NQ = 8;
const int64_t H = s.n_head_kv, D = s.head_dim, NH = s.n_head;
std::mt19937 rng(int8 ? 7 : 11);
std::normal_distribution<float> nd(0.f, 1.f);
Pools ref, slots, host;
ref.alloc(n_blocks, s, int8, false);
slots.alloc(n_slots, s, int8, false);
host.alloc(n_blocks, s, int8, true);
int32_t* ident = dalloc<int32_t>(n_blocks);
{
std::vector<int32_t> t(n_blocks);
for (int64_t i = 0; i < n_blocks; ++i) t[i] = (int32_t) i;
ck(cudaMemcpy(ident, t.data(), n_blocks * 4, cudaMemcpyHostToDevice), "ident");
}
k::KvStreamMap m;
m.page_table = dalloc<int32_t>(n_blocks);
m.slot_block = dalloc<int32_t>(n_slots); m.slot_stamp = dalloc<int32_t>(n_slots); m.slot_ref = dalloc<int32_t>(n_slots);
m.miss_block = dalloc<int32_t>(n_slots); m.miss_slot = dalloc<int32_t>(n_slots);
m.ctl = dalloc<int32_t>(k::kKvCtlInts);
m.n_blocks = n_blocks; m.n_slots = n_slots;
k::kv_stream_reset(m, nullptr);
float* kc = dalloc<float>(H * D);
float* vc = dalloc<float>(H * D);
int32_t* step = dalloc<int32_t>(k::kStepCount);
int32_t* steps = dalloc<int32_t>(NQ * k::kStepCount);
int32_t* ids = dalloc<int32_t>(NQ * cap);
float* q = dalloc<float>(NQ * NH * D);
const uint64_t scr = k::qsa_decode_attn_scratch_floats(cap, s);
float* scratch = dalloc<float>(NQ * scr);
float* out_ref = dalloc<float>(NQ * NH * D);
float* out_str = dalloc<float>(NQ * NH * D);
std::vector<float> hk(H * D), hv(H * D), hq(NQ * NH * D), a((size_t) NQ * NH * D), b2((size_t) NQ * NH * D);
std::vector<int32_t> pt(n_blocks), sb(n_slots);
int batches = 0, bad = 0;
for (int64_t pos = 0; pos < N; ++pos) {
for (auto& x : hk) x = nd(rng) * ((pos % 17 == 0) ? 30.f : 1.f);
for (auto& x : hv) x = nd(rng);
const int32_t st[4] = {(int32_t) pos, (int32_t) (pos + 1), (int32_t) ((pos + 1) / 4), 0};
ck(cudaMemcpy(kc, hk.data(), hk.size() * 4, cudaMemcpyHostToDevice), "k");
ck(cudaMemcpy(vc, hv.data(), hv.size() * 4, cudaMemcpyHostToDevice), "v");
ck(cudaMemcpy(step, st, sizeof(st), cudaMemcpyHostToDevice), "step");
append(ref, ident, step, kc, vc, s, int8, nullptr);
append(slots, m.page_table, step, kc, vc, s, int8, &host.p);
if (pos % 131 != 130 && pos != N - 1) continue;
// a batch: n_q queries at the last n_q positions (as a verify window), each with its own selection
const int n_q = 1 + (int) (rng() % NQ);
std::vector<int32_t> hids((size_t) NQ * cap, 0), hst(NQ * 4, 0);
for (int t = 0; t < n_q; ++t) {
const int64_t p = std::max<int64_t>(0, pos - (n_q - 1 - t)), n_kv = p + 1;
const int64_t width = std::min<int64_t>(n_kv, cap);
const std::vector<int32_t> sel = selection(n_kv, width, rng);
if ((int64_t) sel.size() != width) { std::fprintf(stderr, "selection size %zu != %lld\n", sel.size(), (long long) width); return false; }
std::copy(sel.begin(), sel.end(), hids.begin() + (size_t) t * cap);
hst[t * 4 + 0] = (int32_t) p; hst[t * 4 + 1] = (int32_t) n_kv; hst[t * 4 + 2] = (int32_t) (n_kv / 4);
hst[t * 4 + 3] = (int32_t) width;
}
for (auto& x : hq) x = nd(rng);
ck(cudaMemcpy(ids, hids.data(), hids.size() * 4, cudaMemcpyHostToDevice), "ids");
ck(cudaMemcpy(steps, hst.data(), hst.size() * 4, cudaMemcpyHostToDevice), "steps");
ck(cudaMemcpy(q, hq.data(), hq.size() * 4, cudaMemcpyHostToDevice), "q");
k::qsa_decode_attn_batch(q, ref.attn(ident), ids, steps, cap, s, scratch, out_ref, n_q, nullptr);
k::kv_stream_resolve(m, slots.attn(m.page_table), host.p, int8, ids, steps, n_q, cap, s, nullptr);
k::qsa_decode_attn_batch(q, slots.attn(m.page_table), ids, steps, cap, s, scratch, out_str, n_q, nullptr);
ck(cudaDeviceSynchronize(), "batch");
ck(cudaMemcpy(a.data(), out_ref, (size_t) n_q * NH * D * 4, cudaMemcpyDeviceToHost), "a");
ck(cudaMemcpy(b2.data(), out_str, (size_t) n_q * NH * D * 4, cudaMemcpyDeviceToHost), "b");
if (std::memcmp(a.data(), b2.data(), (size_t) n_q * NH * D * 4) != 0) {
if (bad++ < 5) std::fprintf(stderr, " %s pos %lld n_q %d: streamed attention differs\n", int8 ? "int8" : "fp16", (long long) pos, n_q);
}
// the map inverts itself
ck(cudaMemcpy(pt.data(), m.page_table, n_blocks * 4, cudaMemcpyDeviceToHost), "pt");
ck(cudaMemcpy(sb.data(), m.slot_block, n_slots * 4, cudaMemcpyDeviceToHost), "sb");
int64_t resident = 0;
for (int64_t bk = 0; bk < n_blocks; ++bk) {
if (pt[bk] < -1 || pt[bk] >= n_slots || (pt[bk] >= 0 && sb[pt[bk]] != bk)) {
if (bad++ < 5) std::fprintf(stderr, " map broken at block %lld (table %d)\n", (long long) bk, pt[bk]);
break;
}
resident += pt[bk] >= 0;
}
for (int64_t sl = 0; sl < n_slots; ++sl)
if (sb[sl] >= 0 && pt[sb[sl]] != sl) { if (bad++ < 5) std::fprintf(stderr, " slot %lld not in the table\n", (long long) sl); break; }
++batches;
}
const k::KvStreamCounters c = k::kv_stream_counters(m);
std::printf(" %s: %d batches, %llu block lookups, %llu misses (%.1f%% hit), overflow %d, %d failures\n",
int8 ? "int8" : "fp16", batches, (unsigned long long) c.lookups, (unsigned long long) c.misses,
c.lookups ? 100.0 * (double) (c.lookups - c.misses) / (double) c.lookups : 0.0, (int) c.overflow, bad);
if (c.overflow || c.calls != (uint64_t) batches || c.misses == 0 || c.misses >= c.lookups) ++bad;
// 4. a ring over the same host copy: blocks [b1 - R, b1) restored, read through `b % R`
{
const int64_t R = 1500, b1 = n_blocks, b0 = b1 - R;
Pools ring;
ring.alloc(R, s, int8, false);
int32_t* rt = dalloc<int32_t>(n_blocks);
k::kv_ring_table(rt, n_blocks, R, nullptr);
k::kv_ring_restore(ring.attn(rt), host.p, int8, b0, b1, R, s, nullptr);
const int64_t width = cap;
std::vector<int32_t> hids((size_t) cap), hst = {(int32_t) (N - 1), (int32_t) N, (int32_t) (N / 4), (int32_t) width};
for (int64_t i = 0; i < width; ++i) hids[i] = (int32_t) (N - width + i); // the window's last cells
ck(cudaMemcpy(ids, hids.data(), hids.size() * 4, cudaMemcpyHostToDevice), "ids");
ck(cudaMemcpy(steps, hst.data(), 16, cudaMemcpyHostToDevice), "steps");
k::qsa_decode_attn_batch(q, ref.attn(ident), ids, steps, cap, s, scratch, out_ref, 1, nullptr);
k::qsa_decode_attn_batch(q, ring.attn(rt), ids, steps, cap, s, scratch, out_str, 1, nullptr);
ck(cudaDeviceSynchronize(), "ring");
ck(cudaMemcpy(a.data(), out_ref, (size_t) NH * D * 4, cudaMemcpyDeviceToHost), "a");
ck(cudaMemcpy(b2.data(), out_str, (size_t) NH * D * 4, cudaMemcpyDeviceToHost), "b");
const bool ok = std::memcmp(a.data(), b2.data(), (size_t) NH * D * 4) == 0;
std::printf(" %s ring restore: %s\n", int8 ? "int8" : "fp16", ok ? "identical" : "DIFFERS");
if (!ok) ++bad;
}
return bad == 0;
}
} // namespace
int main() {
std::printf("kv_stream_parity: streamed vs resident KV, bitwise\n");
const bool a = run(true), b = run(false);
if (!a || !b) ++g_fail;
std::printf(g_fail ? "FAIL\n" : "PASS\n");
return g_fail ? 1 : 0;
}
+34 -8
View File
@@ -335,10 +335,12 @@ __global__ void gate_attn_kernel(const float* __restrict__ a, const float* __res
o16[i] = hf(a[i] * (1.0f / (1.0f + expf(-qf[t * 24 * 512 + h * 512 + 256 + d]))));
}
// one block per (token, kv head, 64-value group); 64 threads
// one block per (token, kv head, 64-value group); 64 threads. KV streaming: the pool page only if the block is
// resident (table >= 0), and the host copy and the prompt path's staging pool (both identity layout) when given.
__global__ void kv_append_kernel(const float* __restrict__ K, const float* __restrict__ V, int64_t pos0,
const int32_t* __restrict__ table, int64_t page_size, uint16_t* k_pool,
uint16_t* v_pool, int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale) {
uint16_t* v_pool, int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale,
strata::kernels::KvHostPools host, strata::kernels::KvHostPools stage) {
const int64_t t = blockIdx.x;
const int kvh = blockIdx.y, g = blockIdx.z >> 1;
const bool is_v = (blockIdx.z & 1) != 0;
@@ -347,8 +349,12 @@ __global__ void kv_append_kernel(const float* __restrict__ K, const float* __res
const int64_t pos = pos0 + t;
const int64_t page = table[pos / page_size];
const int64_t row = (page * 2 + kvh) * page_size + pos % page_size;
const int64_t row_id = ((pos / page_size) * 2 + kvh) * page_size + pos % page_size;
if (k_pool != nullptr) {
(is_v ? v_pool : k_pool)[row * 256 + d] = hf(x);
const uint16_t h = hf(x);
if (page >= 0) (is_v ? v_pool : k_pool)[row * 256 + d] = h;
if (host.k_pool != nullptr) (is_v ? host.v_pool : host.k_pool)[row_id * 256 + d] = h;
if (stage.k_pool != nullptr) (is_v ? stage.v_pool : stage.k_pool)[row_id * 256 + d] = h;
return;
}
float a = fabsf(x);
@@ -361,13 +367,27 @@ __global__ void kv_append_kernel(const float* __restrict__ K, const float* __res
const float sf = __half2float(__ushort_as_half(sb));
int q = 0;
if (sf > 0.0f) { q = __float2int_rn(x / sf); q = q < -127 ? -127 : (q > 127 ? 127 : q); }
(is_v ? v_q : k_q)[row * 256 + d] = (int8_t) q;
if (threadIdx.x == 0) (is_v ? v_scale : k_scale)[row * 4 + g] = sb;
if (page >= 0) {
(is_v ? v_q : k_q)[row * 256 + d] = (int8_t) q;
if (threadIdx.x == 0) (is_v ? v_scale : k_scale)[row * 4 + g] = sb;
}
if (host.k_q != nullptr) {
(is_v ? host.v_q : host.k_q)[row_id * 256 + d] = (int8_t) q;
if (threadIdx.x == 0) (is_v ? host.v_scale : host.k_scale)[row_id * 4 + g] = sb;
}
if (stage.k_q != nullptr) {
(is_v ? stage.v_q : stage.k_q)[row_id * 256 + d] = (int8_t) q;
if (threadIdx.x == 0) (is_v ? stage.v_scale : stage.k_scale)[row_id * 4 + g] = sb;
}
}
__global__ void to_f16_kernel(const float* __restrict__ x, uint16_t* __restrict__ y, int64_t n) {
for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x * blockDim.x)
y[i] = hf(x[i]);
}
__global__ void round_f16_kernel(const float* __restrict__ x, float* __restrict__ y, int64_t n) {
for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x * blockDim.x)
y[i] = __half2float(__float2half_rn(x[i]));
}
__global__ void to_bf16_kernel(const float* __restrict__ x, uint16_t* __restrict__ y, int64_t n) {
for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x * blockDim.x)
y[i] = bf(x[i]);
@@ -377,10 +397,11 @@ __global__ void to_bf16_kernel(const float* __restrict__ x, uint16_t* __restrict
void kv_append(const float* K, const float* V, int64_t T, int64_t pos0, const int32_t* page_table, int64_t page_size,
uint16_t* k_pool, uint16_t* v_pool, int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale,
void* stream) {
void* stream, const strata::kernels::KvHostPools* host, const strata::kernels::KvHostPools* stage) {
if (T <= 0) return;
kv_append_kernel<<<dim3((unsigned) T, 2, 8), 64, 0, (cudaStream_t) stream>>>(K, V, pos0, page_table, page_size, k_pool,
v_pool, k_q, v_q, k_scale, v_scale);
kv_append_kernel<<<dim3((unsigned) T, 2, 8), 64, 0, (cudaStream_t) stream>>>(
K, V, pos0, page_table, page_size, k_pool, v_pool, k_q, v_q, k_scale, v_scale,
host ? *host : strata::kernels::KvHostPools{}, stage ? *stage : strata::kernels::KvHostPools{});
check("kv_append");
}
void to_f16(const float* x, uint16_t* y, int64_t n, void* stream) {
@@ -388,6 +409,11 @@ void to_f16(const float* x, uint16_t* y, int64_t n, void* stream) {
to_f16_kernel<<<(unsigned) ((n + 255) / 256 < 4096 ? (n + 255) / 256 : 4096), 256, 0, (cudaStream_t) stream>>>(x, y, n);
check("to_f16");
}
void round_f16(const float* x, float* y, int64_t n, void* stream) {
if (n <= 0) return;
round_f16_kernel<<<(unsigned) ((n + 255) / 256 < 4096 ? (n + 255) / 256 : 4096), 256, 0, (cudaStream_t) stream>>>(x, y, n);
check("round_f16");
}
void to_bf16(const float* x, uint16_t* y, int64_t n, void* stream) {
if (n <= 0) return;
to_bf16_kernel<<<(unsigned) ((n + 255) / 256 < 4096 ? (n + 255) / 256 : 4096), 256, 0, (cudaStream_t) stream>>>(x, y, n);
+142 -5
View File
@@ -9,6 +9,8 @@
#include "strata/kernels/iq_kernels.hpp"
#include "strata/kernels/cpu/expert_layout.hpp"
#include "strata/kernels/qsa.hpp"
#include "strata/kernels/kv_stream.hpp"
#include "strata/core/layer.hpp"
#include "strata/kernels/qsa_decode_attn.hpp"
#include "strata/kernels/qsa_select.hpp"
#include "strata/prefill/gemm.hpp"
@@ -111,8 +113,44 @@ struct Prefill::Impl {
std::vector<uint32_t> ple_rows;
float* ple_norm = nullptr;
PrefillStats* stats = nullptr;
// KV streaming: one layer's whole K/V, staged from the host copy per layer and chunk (identity layout)
strata::kernels::KvHostPools stage;
int32_t* ident_table = nullptr;
};
namespace {
// the staging pool of a streamed session: every page of one layer (same sequence in init and bytes_needed)
// STRATA_KV_STAGE_OWN (A/B only): the staging pool gets its own allocation instead of borrowed expert slots, so a
// streamed run lends the prompt path exactly the slots a resident one does (a lent expert runs on the CPU, which
// rounds differently: without this an A/B compares two expert placements as well as two KV placements)
bool stage_own() { static const bool v = std::getenv("STRATA_KV_STAGE_OWN") != nullptr; return v; }
void take_stage(Alloc& o_borrowed, const core::SessionState& ss, const strata::kernels::QsaShapes& s,
strata::kernels::KvHostPools& st, bool& ok) {
const core::QsaState& q0 = ss.qsa_states[0];
if (q0.kv_mode != 1) return;
if (stage_own() && o_borrowed.count_only) return;
Alloc own;
own.owned = o_borrowed.owned;
Alloc& o = stage_own() ? own : o_borrowed;
const size_t rows = (size_t) q0.n_pages * s.n_head_kv * s.page_size;
if (q0.kv_int8) {
st.k_q = o.take<int8_t>(rows * s.head_dim, ok);
st.v_q = o.take<int8_t>(rows * s.head_dim, ok);
st.k_scale = o.take<uint16_t>(rows * (s.head_dim / 64), ok);
st.v_scale = o.take<uint16_t>(rows * (s.head_dim / 64), ok);
} else {
st.k_pool = o.take<uint16_t>(rows * s.head_dim, ok);
st.v_pool = o.take<uint16_t>(rows * s.head_dim, ok);
}
}
strata::kernels::QsaAttnPools pools_of(const strata::kernels::KvHostPools& h, const int32_t* table) {
strata::kernels::QsaAttnPools p;
p.k_pool = h.k_pool; p.v_pool = h.v_pool; p.k_q = h.k_q; p.v_q = h.v_q; p.k_scale = h.k_scale; p.v_scale = h.v_scale;
p.page_table = table;
return p;
}
} // namespace
Prefill::Prefill() : impl_(new Impl) {}
Prefill::~Prefill() {
if (!impl_) return;
@@ -189,6 +227,17 @@ bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, co
}
m.ple_emb = o.take<float>(T * N, ok);
m.ple_norm = o.take<float>((size_t) strata::kernels::NG_HC_DIM, ok);
take_stage(o, ss, s, m.stage, ok);
if (ok && m.stage.present()) {
const int64_t pages = ss.qsa_states[0].n_pages;
std::vector<int32_t> ident((size_t) pages);
for (int64_t i = 0; i < pages; ++i) ident[(size_t) i] = (int32_t) i;
if (cudaMalloc((void**) &m.ident_table, ident.size() * 4) != cudaSuccess ||
cudaMemcpy(m.ident_table, ident.data(), ident.size() * 4, cudaMemcpyHostToDevice) != cudaSuccess)
ok = false;
else
m.owned.push_back(m.ident_table);
}
m.ids_host.resize(T * K); m.slot_host.resize(T * K); m.src_host.resize(T * K); m.cnt.resize(NE); m.off.resize(NE + 1);
m.ple_emb_host.resize(T * N); m.ple_rows.resize(T * strata::kernels::PLE_N_HEADS);
if (!ok) { err = "prefill: device buffers for a chunk of " + std::to_string(chunk) + " tokens do not fit"; return false; }
@@ -223,6 +272,8 @@ uint64_t Prefill::bytes_needed(const core::ModelGeometry& g, const core::Session
for (int i = 0; i < STAGE; ++i) o.take<uint8_t>((size_t) MAXBLOB(), ok);
f(T * N);
f((size_t) strata::kernels::NG_HC_DIM);
strata::kernels::KvHostPools stage;
take_stage(o, ss, s, stage, ok);
return o.used + (8u << 20); // alignment slack
}
@@ -261,6 +312,8 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]};
for (int64_t c0 = 0; c0 < n; c0 += m.T) {
if (should_stop && should_stop()) { err = "cancelled"; return false; }
if (std::getenv("STRATA_TRACE")) { std::fprintf(stderr, "strata trace: prompt chunk %lld of %lld\n", (long long) c0, (long long) n); std::fflush(stderr); }
const int64_t T = std::min(m.T, n - c0), p0 = pos0 + c0;
++stats_.chunks;
// ---- embeddings, broadcast to the four streams
@@ -367,8 +420,15 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
if (!bf16_proj(m.gemm, wiq, m.mixed_bf, m.q_idx, T, v.name("indexer.q_proj.weight"), err)) return false;
rms_rows(m.Kc, (const float*) wkn->data, T * 2, 256, 256, EPS, m.cs);
rope(m.Kc, T, 2, 256, 512, p0, (float) strata::kernels::qsa_freq_base(), m.cs);
// KV streaming: this layer's cells [0, p0) come in from the host copy to the staging pool, and the
// chunk's cells go to the host copy, the staging pool, and the VRAM slots of resident blocks
const bool staged = st.kv_mode == 1;
if (staged)
strata::kernels::kv_stage_from_host(pools_of(m.stage, m.ident_table), st.host, st.kv_int8,
(p0 + s.page_size - 1) / s.page_size, s, m.cs);
kv_append(m.Kc, m.Vc, T, p0, st.page_table, s.page_size, st.kv_int8 ? nullptr : st.k_pool,
st.kv_int8 ? nullptr : st.v_pool, st.k_q, st.v_q, st.k_scale, st.v_scale, m.cs);
st.kv_int8 ? nullptr : st.v_pool, st.k_q, st.v_q, st.k_scale, st.v_scale, m.cs,
&st.host, staged ? &m.stage : nullptr);
split_q(m.Qf, m.q, T, m.cs);
rms_rows(m.q, (const float*) wqn->data, T * 24, 256, 256, EPS, m.cs);
rope(m.q, T, 24, 256, 6144, p0, (float) strata::kernels::qsa_freq_base(), m.cs);
@@ -393,10 +453,73 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
strata::kernels::qsa_block_topk(m.sel_scores, steps0, nb, m.max_blocks, m.cap, s,
m.sel_ids + t0 * m.cap, m.cs);
}
strata::kernels::QsaAttnPools pools;
pools.page_table = st.page_table;
if (st.kv_int8) { pools.k_q = st.k_q; pools.v_q = st.v_q; pools.k_scale = st.k_scale; pools.v_scale = st.v_scale; }
else { pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; }
// STRATA_IDX_FP16_CHECK: would FP16 pooled indexer keys select the same cells? (the KV-streaming
// design's last question). Every query is selected again from the pooled keys and `dead` rounded
// to fp16 (exactly what an fp16 store reads back); the agreement with the fp32 selection is
// printed cumulatively after each chunk's last QSA layer. Debug: syncs per layer.
if (static const bool f16chk = std::getenv("STRATA_IDX_FP16_CHECK") != nullptr; f16chk) {
static float *pooled16 = nullptr, *dead16 = nullptr;
static int32_t* ids16 = nullptr;
static double shared = 0, cells = 0;
static long long queries = 0, same = 0, sel_queries = 0;
const int64_t rows = st.idx_pooled_rows;
if (pooled16 == nullptr &&
(cudaMalloc((void**) &pooled16, (size_t) rows * s.idx_dim * 4) != cudaSuccess ||
cudaMalloc((void**) &dead16, (size_t) s.idx_dim * 4) != cudaSuccess ||
cudaMalloc((void**) &ids16, (size_t) (m.T * m.cap) * 4) != cudaSuccess)) {
err = "STRATA_IDX_FP16_CHECK: no room for its buffers";
return false;
}
round_f16(st.idx_pooled, pooled16, rows * s.idx_dim, m.cs);
round_f16(st.idx_dead, dead16, s.idx_dim, m.cs);
for (int64_t t0 = 0; t0 < T; t0 += m.sel_batch) {
const int64_t nb = std::min(m.sel_batch, T - t0);
const int32_t* steps0 = m.steps_dev + t0 * strata::kernels::kStepCount;
strata::kernels::qsa_block_scores(pooled16, dead16, m.q_idx + t0 * 512, steps0, nb,
m.max_blocks, s, m.sel_scores, m.cs);
strata::kernels::qsa_block_topk(m.sel_scores, steps0, nb, m.max_blocks, m.cap, s,
ids16 + t0 * m.cap, m.cs);
}
std::vector<int32_t> a((size_t) (T * m.cap)), b((size_t) (T * m.cap));
cudaMemcpyAsync(a.data(), m.sel_ids, a.size() * 4, cudaMemcpyDeviceToHost, m.cs);
cudaMemcpyAsync(b.data(), ids16, b.size() * 4, cudaMemcpyDeviceToHost, m.cs);
cudaStreamSynchronize(m.cs);
for (int64_t t = 0; t < T; ++t) {
const int64_t w = m.steps_host[(size_t) (t * strata::kernels::kStepCount + strata::kernels::kStepWidth)];
const int32_t *x = a.data() + t * m.cap, *y = b.data() + t * m.cap;
int64_t i = 0, j = 0, c = 0;
while (i < w && j < w) {
if (x[i] == y[j]) { ++c; ++i; ++j; } else if (x[i] < y[j]) ++i; else ++j;
}
++queries;
same += c == w;
if (p0 + t + 1 > m.cap) { ++sel_queries; shared += (double) c; cells += (double) w; }
}
if (qsa_index + 1 == g.n_qsa_layers())
std::fprintf(stderr, "strata prefill: FP16 indexer keys: %lld of %lld selections identical; "
"where the selection is sparse, %.4f%% of cells shared (%lld queries)\n",
same, queries, cells > 0 ? 100.0 * shared / cells : 100.0, sel_queries);
}
// STRATA_QSA_DUMP=<file>: append every QSA layer's selected cells for the prompt's last
// STRATA_QSA_DUMP_LAST (4096) positions - records of int32 {qsa layer, pos0, T, cap} + T*cap cells,
// for tools/qsa_locality.py (how local the sparse attention's reads are: the KV-streaming question)
if (static const char* dump = std::getenv("STRATA_QSA_DUMP"); dump != nullptr) {
static const long long last = std::getenv("STRATA_QSA_DUMP_LAST")
? std::atoll(std::getenv("STRATA_QSA_DUMP_LAST")) : 4096;
if (p0 + T > pos0 + n - last) {
std::vector<int32_t> h((size_t) (T * m.cap));
cudaMemcpyAsync(h.data(), m.sel_ids, h.size() * 4, cudaMemcpyDeviceToHost, m.cs);
cudaStreamSynchronize(m.cs);
if (std::FILE* f = std::fopen(dump, "ab")) {
const int32_t hdr[4] = {(int32_t) qsa_index, (int32_t) p0, (int32_t) T, (int32_t) m.cap};
std::fwrite(hdr, 4, 4, f);
std::fwrite(h.data(), 4, h.size(), f);
std::fclose(f);
}
}
}
const strata::kernels::QsaAttnPools pools = staged ? pools_of(m.stage, m.ident_table)
: core::qsa_attn_pools(st);
for (int64_t t0 = 0; t0 < T; t0 += m.attn_batch) {
const int64_t nb = std::min(m.attn_batch, T - t0);
strata::kernels::qsa_decode_attn_batch(m.q + t0 * ZV, pools, m.sel_ids + t0 * m.cap,
@@ -531,6 +654,20 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
return false;
}
stats_.ms_total += ms_since(t_start);
if (std::getenv("STRATA_STATE_HASH_GDN") != nullptr) { // debug: the GDN states as the prompt path leaves them
cudaStreamSynchronize(m.cs);
std::vector<uint8_t> b((size_t) gdn_floats * 4);
std::string line;
char h[8];
for (int64_t i = 0; i < g.n_gdn_layers(); ++i) {
cudaMemcpy(b.data(), ss.gdn_state + (size_t) i * gdn_floats, b.size(), cudaMemcpyDeviceToHost);
uint64_t x = 1469598103934665603ull;
for (uint8_t c : b) x = (x ^ c) * 1099511628211ull;
std::snprintf(h, sizeof(h), "%04llx ", (unsigned long long) (x & 0xffff));
line += h;
}
std::fprintf(stderr, "strata prefill: GDN_HASH %s\n", line.c_str());
}
return true;
}
+639 -47
View File
@@ -41,6 +41,7 @@
#include "strata/prefill/prefill.hpp"
#include "strata/core/native_dense.hpp"
#include "strata/program/logits_selection.hpp"
#include "strata/spec/suffix_drafter.hpp"
#include <cuda_runtime.h>
@@ -48,6 +49,10 @@
#include <algorithm>
#include <iostream>
#include <thread>
#include <atomic>
#include <condition_variable>
#include <deque>
#include <mutex>
#include <charconv>
#include <cmath>
#include <cstdio>
@@ -112,6 +117,7 @@ struct Options {
double ple_delay_us = 0; ///< fault injection: every row read completes no earlier than this
bool ple_sync_submit = false; ///< A/B arm: submit reads on the token thread, no I/O worker
std::string kv = "fp16"; ///< plan v0.3 P7: KV storage, fp16 (default) or int8 (half the VRAM)
int64_t kv_resident = 0; ///< KV streaming: resident cells per QSA layer (0: all in VRAM)
std::string dump_residual;
/// The head input, `bb.mixed`. It exists so the head can be SPLIT: steps 1-4 (the per-stream norm, the two
/// bf16 projections and the stream mean) recompute cheaply in Python, and only the 794 MB GEMV does not.
@@ -231,6 +237,25 @@ struct Options {
/// The vision path: keep a per-cell (t, h, w) rotary position table so --serve can take GENI requests.
bool vision = false;
int adapt_swaps = 96;
/// --serve: how many conversation checkpoints to keep between requests (0 = every request reads its whole
/// prompt again, the v0.1.2 behaviour). One is the GDN recurrence of the 36 layers, the QSA indexer tails and
/// the PLE history (~118 MB of host RAM); the KV cache itself is positional and stays where it is.
int prompt_cache = 6;
/// --serve: also keep a checkpoint every N freshly read prompt tokens (0 = only at the last turn boundary)
int64_t prompt_cache_every = 16384;
/// --serve: the token that opens a chat turn (<|im_start|>). The last one in a prompt is where the chat's
/// history ends and the new assistant turn begins, which is the checkpoint the next request can reuse.
int64_t turn_token = 248045;
/// --serve: a text part of the prompt of at most N tokens (a chat message, the assistant header, a short tool
/// result) goes through the verify windows, S tokens at a time, instead of the batched prompt path (0 = always
/// the batched path)
int64_t short_read = 64;
/// The suffix drafter: when the text being written repeats an earlier stretch of the context (code edits, quoted
/// input, tool-call JSON) by at least this many tokens, the window is filled with what followed it there instead
/// of the MTP's drafts, up to --spec - 1 of them, as long as the MTP's own first guess agrees. 0 = MTP only.
int suffix_draft = 0;
/// The MTP's own window cap (0 = --spec): with --spec 6 --mtp-max-t 4 the long windows come from suffix matches.
int mtp_max_t = 0;
};
void usage() {
@@ -250,6 +275,9 @@ void usage() {
" --ple-sync-submit A/B arm: submit table reads on the token thread (default: an I/O thread)\n"
" --kv fp16|int8 KV storage (plan v0.3 P7): int8 codes + fp16 scale per 64 values, half the\n"
" VRAM; default fp16 until gate G-C accepts int8\n"
" --kv-resident N KV streaming: keep N cells of each QSA layer in VRAM (min 20480) and the\n"
" whole K/V in pinned RAM; the freed VRAM goes to expert slots. 0 (default):\n"
" all of it in VRAM. A context of N cells or fewer is not streamed\n"
" --stream-token enqueue token work on the session stream (experimental)\n"
" --check-logits copy and check all logits in the stream-token path\n"
" --gr-fp32-activations experimental CUDA-oracle GR activation precision\n"
@@ -288,6 +316,14 @@ void usage() {
" --shared-late A/B: shared expert after the CPU pool (default: overlapped with it)\n"
" --keep-canonical A/B: also load canonical copies of natively served tensors (more VRAM)\n"
" --vision --serve takes images too (GENI requests; embeddings from strata-vision)\n"
" --prompt-cache N --serve: keep N conversation checkpoints between requests (default 6, ~118 MB\n"
" of RAM each; 0 = read every prompt from the start)\n"
" --prompt-cache-every N --serve: also checkpoint every N fresh prompt tokens (default 16384, 0 = off)\n"
" --turn-token ID --serve: the token that opens a chat turn (default 248045, <|im_start|>)\n"
" --short-read N --serve: read at most N fresh text tokens through the decode windows instead\n"
" of the batched prompt path (default 64, 0 = off)\n"
" --suffix-draft N draft from an earlier repeat of the last N+ tokens of context (0 = MTP only)\n"
" --mtp-max-t M cap the MTP's windows at M tokens (0 = --spec; longer ones come from suffixes)\n"
" --no-token-graph A/B: two graphs per layer (the host launches each) instead of one per token\n"
" --no-fused-gr A/B: the six-kernel hyper-connection read and a separate write (native)\n"
" --prefill CHUNK batched prompt processing in chunks of CHUNK tokens (needs --native)\n"
@@ -422,6 +458,15 @@ void drive_pool_multi(void* user, const float* x_f, const int32_t* ids, int64_t
++t->calls;
}
/// STRATA_TRACE=1: the VRAM left at a step of the startup (finds what fills the card after the cache is sized)
void mem_mark(const char* where) {
static const bool on = std::getenv("STRATA_TRACE") != nullptr;
if (!on) return;
size_t free_b = 0, total_b = 0;
cudaMemGetInfo(&free_b, &total_b);
std::fprintf(stderr, "strata trace: %lld MiB free after %s\n", (long long) (free_b >> 20), where);
}
int argmax(const std::vector<float>& v) {
int best = 0;
for (size_t i = 1; i < v.size(); ++i)
@@ -429,6 +474,84 @@ int argmax(const std::vector<float>& v) {
return best;
}
// ---- --serve's conversation cache. A chat or an agent sends the whole conversation again with every request, and
// reading it again is what made a long session wait minutes for every turn. What a sequence leaves behind splits in
// two, and only one half needs copying:
// * POSITIONAL state - the KV cache of the 12 QSA layers and their pooled indexer keys, the draft layer's KV. A
// cell is written once for its position and read only by later positions (the block scores take `dead` for the
// block being filled, never its pooled row), so rewinding to a position just means writing from there again.
// * RUNNING state - the 36 GDN recurrences and conv histories, each QSA layer's indexer tail (the unfinished
// block's raw keys) and the PLE's normalized history. Each describes "everything so far" and cannot be
// rewound, so a checkpoint is a copy of exactly these: ~118 MB, the same set the verifier snapshots to roll
// back rejected drafts.
// A checkpoint is only valid while the positional cells below it still hold ITS tokens, so the serve loop keeps
// just the checkpoints that are prefixes of the tokens the session holds now.
struct ImgKey {
int64_t start = 0; ///< the image's first <|image_pad|> position
uint64_t hash = 0; ///< its embeddings and grid: the pad tokens alone are the same for every picture
bool operator==(const ImgKey& o) const { return start == o.start && hash == o.hash; }
};
struct ConvCheckpoint {
std::vector<int32_t> ids; ///< the tokens this state has consumed
std::vector<ImgKey> imgs; ///< the images among them
std::vector<uint8_t> gdn, ple, tails;
};
uint64_t fnv1a(const void* data, size_t n, uint64_t h = 1469598103934665603ull) {
const uint8_t* p = (const uint8_t*) data;
for (size_t i = 0; i < n; ++i) { h ^= p[i]; h *= 1099511628211ull; }
return h;
}
struct ConvStateSizes {
size_t gdn = 0, ple = 0, tail = 0;
};
ConvStateSizes conv_state_sizes(const strata::core::ModelGeometry& g) {
ConvStateSizes z;
z.gdn = (size_t) g.n_gdn_layers() *
((size_t) g.ssm_state_size * (size_t) g.ssm_v_heads * (size_t) g.ssm_state_size +
(size_t) g.ssm_conv_channels * (size_t) (g.ssm_d_conv - 1)) * sizeof(float);
z.ple = (size_t) strata::kernels::NG_HIST * (size_t) strata::kernels::NG_HC_DIM * sizeof(float);
z.tail = (size_t) (strata::kernels::qsa_real_shapes().idx_block - 1) * (size_t) g.idx_key_dim * sizeof(float);
return z;
}
/// Copies the running state out. The caller has synchronized the device.
bool checkpoint_save(ConvCheckpoint& c, const strata::core::SessionState& ss, const strata::core::ModelGeometry& g) {
const ConvStateSizes z = conv_state_sizes(g);
c.gdn.resize(z.gdn);
c.ple.resize(ss.ple_hist != nullptr ? z.ple : 0);
c.tails.resize(z.tail * (size_t) g.n_qsa_layers());
if (cudaMemcpy(c.gdn.data(), ss.gdn_state, z.gdn, cudaMemcpyDeviceToHost) != cudaSuccess) return false;
if (!c.ple.empty() && cudaMemcpy(c.ple.data(), ss.ple_hist, z.ple, cudaMemcpyDeviceToHost) != cudaSuccess)
return false;
for (int64_t i = 0; i < g.n_qsa_layers(); ++i)
if (cudaMemcpy(c.tails.data() + (size_t) i * z.tail, ss.qsa_states[i].idx_tail, z.tail, cudaMemcpyDeviceToHost) !=
cudaSuccess)
return false;
return true;
}
/// Puts a checkpoint's running state back; the positional cells below it are the caller's to guarantee.
bool checkpoint_restore(const ConvCheckpoint& c, strata::core::SessionState& ss, const strata::core::ModelGeometry& g) {
const ConvStateSizes z = conv_state_sizes(g);
if (c.gdn.size() != z.gdn || c.tails.size() != z.tail * (size_t) g.n_qsa_layers()) return false;
if (cudaMemcpy(ss.gdn_state, c.gdn.data(), z.gdn, cudaMemcpyHostToDevice) != cudaSuccess) return false;
if (!c.ple.empty() && cudaMemcpy(ss.ple_hist, c.ple.data(), z.ple, cudaMemcpyHostToDevice) != cudaSuccess)
return false;
for (int64_t i = 0; i < g.n_qsa_layers(); ++i)
if (cudaMemcpy(ss.qsa_states[i].idx_tail, c.tails.data() + (size_t) i * z.tail, z.tail, cudaMemcpyHostToDevice) !=
cudaSuccess)
return false;
// the PLE's token window is the last two tokens, OLDEST FIRST, -1 where there is none (as session_zero leaves it)
const size_t L = c.ids.size();
ss.ple_prev[0] = L >= 2 ? c.ids[L - 2] : -1;
ss.ple_prev[1] = L >= 1 ? c.ids[L - 1] : -1;
return cudaDeviceSynchronize() == cudaSuccess;
}
} // namespace
int main(int argc, char** argv) {
@@ -493,6 +616,7 @@ int main(int argc, char** argv) {
else if (a == "--ple-delay-us") o.ple_delay_us = std::atof(next("--ple-delay-us"));
else if (a == "--ple-sync-submit") o.ple_sync_submit = true;
else if (a == "--kv") o.kv = next("--kv");
else if (a == "--kv-resident") o.kv_resident = std::atoll(next("--kv-resident"));
else if (a == "--stream-token") o.stream_token = true;
else if (a == "--check-logits") o.check_logits = true;
else if (a == "--gr-fp32-activations") o.gr_fp32_activations = true;
@@ -544,6 +668,12 @@ int main(int argc, char** argv) {
else if (a == "--pcie-mode") o.pcie_mode = next("--pcie-mode");
else if (a == "--serve") o.serve = true;
else if (a == "--vision") o.vision = true;
else if (a == "--prompt-cache") o.prompt_cache = std::max(0, std::atoi(next("--prompt-cache")));
else if (a == "--prompt-cache-every") o.prompt_cache_every = std::max(0LL, std::atoll(next("--prompt-cache-every")));
else if (a == "--turn-token") o.turn_token = std::atoll(next("--turn-token"));
else if (a == "--short-read") o.short_read = std::max(0LL, std::atoll(next("--short-read")));
else if (a == "--suffix-draft") o.suffix_draft = std::max(0, std::atoi(next("--suffix-draft")));
else if (a == "--mtp-max-t") o.mtp_max_t = std::max(0, std::atoi(next("--mtp-max-t")));
else if (a == "--no-spec-split") o.spec_split = false;
else if (a == "--eos-ids") {
std::string e;
@@ -597,6 +727,11 @@ int main(int argc, char** argv) {
return 2;
}
strata::core::qsa_set_kv_int8(o.kv == "int8");
if (o.kv_resident < 0) {
std::fprintf(stderr, "strata generate: --kv-resident must be >= 0\n");
return 2;
}
strata::core::qsa_set_kv_resident(o.kv_resident);
strata::core::layer_set_shared_early(!o.shared_late);
if (!o.native_preset.empty()) {
if (o.no_ple || o.ple_gguf.empty()) {
@@ -827,6 +962,10 @@ int main(int argc, char** argv) {
std::fprintf(stderr, "strata generate: session_init failed\n");
return 1;
}
if (g.n_qsa_layers() > 0 && ss.qsa_states[0].kv_mode == 1)
std::fprintf(stderr, "strata generate: KV streaming: %lld of %lld cells per QSA layer in VRAM, the K/V in "
"%.2f GiB of pinned RAM\n", (long long) (ss.qsa_states[0].n_slots * 4),
(long long) o.max_context, (double) strata::core::qsa_kv_host_bytes() / 1073741824.0);
// ---- **THE HALF-LEVEL DUMP HAS TO BE ARMED BEFORE `session_capture`, AND THE LADDER MUST NOT BE.** The
// half copies are issued from inside `block_layer_pre`/`block_layer_post`, so they are only ever enqueued
@@ -1003,6 +1142,7 @@ int main(int argc, char** argv) {
// ---- R4's slot storage. Allocated AFTER the weights and the session, so `cudaMemGetInfo` inside `open`
// sees the memory this process actually has left rather than the card's idle figure - and refuses with both
// numbers if the slots do not fit, instead of handing back a cache smaller than it was asked for.
mem_mark("the weights, the session and the drafter");
strata::core::ExpertCache xcache;
std::vector<std::pair<int32_t, int32_t>> profile;
if (!o.expert_profile.empty()) {
@@ -1018,6 +1158,7 @@ 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);
}
const bool auto_cache = o.expert_cache < 0;
if (o.expert_cache < 0) {
size_t free_b = 0, total_b = 0;
cudaMemGetInfo(&free_b, &total_b);
@@ -1053,15 +1194,51 @@ int main(int argc, char** argv) {
o.expert_cache = (int) sized_slots.size();
}
if (o.expert_cache > 0) {
const bool ok = sized_slots.empty()
? xcache.open(o.expert_cache, g.n_layers, g.n_expert, (int64_t) strata::kernels::cpu::expert_layout().max_blob, err)
: xcache.open_sized(sized_slots, g.n_layers, g.n_expert, err);
if (!ok) {
std::fprintf(stderr, "strata generate: %s\n", err.c_str());
return 1;
// With `--expert-cache auto` the reserve must still be free once the slots are WRITTEN: under WDDM an
// allocation is not resident until it is touched, and the free figure read before it can be ~1 GB too
// high. A cache sized from it filled the card to 0 MiB, the driver then paged, and a request that needed a
// page back while the verify graph spun on a host flag never finished. So the slots are zeroed and the
// free figure read again; while it is short of the reserve the cache is reopened smaller.
for (int attempt = 0;; ++attempt) {
const bool ok = sized_slots.empty()
? xcache.open(o.expert_cache, g.n_layers, g.n_expert, (int64_t) strata::kernels::cpu::expert_layout().max_blob, err)
: xcache.open_sized(sized_slots, g.n_layers, g.n_expert, err);
if (!ok) {
std::fprintf(stderr, "strata generate: %s\n", err.c_str());
return 1;
}
if (!auto_cache || attempt >= 6) break;
cudaMemset(xcache.device_slot(0), 0, (size_t) xcache.bytes());
cudaDeviceSynchronize();
size_t free_b = 0, total_b = 0;
cudaMemGetInfo(&free_b, &total_b);
const int64_t want = (int64_t) o.vram_reserve_mib << 20;
if ((int64_t) free_b >= want - (64ll << 20)) break;
// short by (want - free); a figure of 0 only says "at least", so then give back a quarter as well
int64_t give = want - (int64_t) free_b + (64ll << 20);
if (free_b < ((size_t) 16 << 20)) give = std::max<int64_t>(give, xcache.bytes() / 4);
const int64_t keep_bytes = xcache.bytes() - give;
std::fprintf(stderr, "strata generate: only %lld MiB free once the slots are written (reserve %d MiB); "
"shrinking the expert cache\n", (long long) (free_b >> 20), o.vram_reserve_mib);
xcache.close();
if (keep_bytes <= 0) { o.expert_cache = 0; sized_slots.clear(); break; }
if (!sized_slots.empty()) {
int64_t used = 0;
size_t keep = 0;
while (keep < sized_slots.size() && used + (sized_slots[keep] + 255) / 256 * 256 <= keep_bytes)
used += (sized_slots[keep++] + 255) / 256 * 256;
sized_slots.resize(keep);
o.expert_cache = (int) keep;
} else {
o.expert_cache = (int) (keep_bytes / (int64_t) strata::kernels::cpu::expert_layout().max_blob);
}
if (o.expert_cache <= 0) { o.expert_cache = 0; sized_slots.clear(); break; }
}
}
if (o.expert_cache > 0) {
std::fprintf(stderr, "strata generate: expert cache %lld slots, %.2f GiB of VRAM; policy is\n",
(long long) xcache.slots(), xcache.gib());
mem_mark("opening the expert cache");
xcache.set_per_layer_admission(o.expert_cache_per_layer);
// **ROUND 328: THE HIT PATH IS PROVABLY WRONG, AND THIS SAYS SO OUT LOUD RATHER THAN LETTING IT
// CORRUPT A RUN QUIETLY.** With the cache on, the generated tokens DIVERGE from the cache-off run:
@@ -1109,12 +1286,13 @@ int main(int argc, char** argv) {
// **AND ONE SLOT IS READ BACK AND COMPARED.** A residency table that is right about indices and wrong
// about bytes produces a plausible token, which is this project's most expensive failure mode; the
// cache's own `verify_slot` is the check and it costs one 1.38 MB D2H at startup.
if (!xcache.verify_slot(xcache.slot_of(profile[0].first, profile[0].second),
if (prefilled > 0 && !xcache.verify_slot(xcache.slot_of(profile[0].first, profile[0].second),
srcp->blob(profile[0].first, profile[0].second), err,
(int64_t) strata::kernels::cpu::expert_layout().blob_bytes(profile[0].first))) {
std::fprintf(stderr, "strata generate: %s\n", err.c_str());
return 1;
}
mem_mark("the profile fill");
std::fprintf(stderr, "strata generate: pre-filled %lld of %lld slots from the profile; slot 0 verified\n",
(long long) prefilled, (long long) want);
}
@@ -1172,6 +1350,7 @@ int main(int argc, char** argv) {
drive.d.hit_done = (void*) hit_done;
drive.d.hit_poke = !o.no_hit_poke;
drive.d.h_dst.resize((size_t) K);
mem_mark("the R4 hit path");
std::fprintf(stderr, "strata generate: R4 hit path ON - resident experts are computed on the GPU\n");
}
// ---- P0.S8: the routing trace. Only meaningful with the pool running, because the ids arrive through
@@ -1289,6 +1468,7 @@ int main(int argc, char** argv) {
return 2;
}
mem_mark("the expert cache and the graphs");
std::fprintf(stderr, "strata generate: session is up; locating the head\n");
const strata::core::WeightRef* wo = wt.find("output.weight");
if (wo == nullptr) { std::fprintf(stderr, "strata generate: output.weight is missing\n"); return 1; }
@@ -1758,11 +1938,15 @@ 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>
// DONE <generated> <prompt_tokens> <prompt_ms> <decode_ms> <stop|length|cancel> <drafts accepted>
// <drafts offered> <prompt tokens reused>
//
// (`ERR <message>` instead when a request cannot run; `QUIT` ends the process). Every request starts from an
// empty sequence (`session_zero`): the prompt goes through the batched prompt path and its last token through
// the first verify window - the path all three model files share. Decoding is greedy.
// 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
// cannot run; `STOP` ends the running request at its next step; `QUIT` ends the process.) A request continues
// from the live session or the longest conversation checkpoint its prompt starts with (see ConvCheckpoint),
// otherwise from an empty sequence (`session_zero`); the rest of the prompt goes through the batched prompt path
// and its last token through the first verify window - the path all three model files share. Decoding is greedy.
if (o.serve) {
if (o.spec < 2 || o.mtp.empty() || o.prefill_chunk <= 0 || thits.d_res == nullptr || host_res.empty()) {
std::fprintf(stderr, "strata serve: needs --spec T, --mtp DIR, --prefill CHUNK, --expert-profile P and "
@@ -1773,8 +1957,11 @@ int main(int argc, char** argv) {
void* borrow = nullptr;
uint64_t borrow_bytes = 0;
int32_t lend_first = -1;
if (!o.no_prefill_borrow && d_res != nullptr) {
const uint64_t need = strata::prefill::Prefill::bytes_needed(g, ss, o.prefill_chunk);
// a cache too small to lend the prompt path its buffers would make it allocate them on top - on a card the
// cache already filled to its reserve, that is the over-subscription the auto sizing avoids - so the
// prompt chunk is halved until its buffers fit in the lendable slots (a smaller chunk only reads slower)
for (int64_t chunk = o.prefill_chunk; !o.no_prefill_borrow && d_res != nullptr && chunk >= 256; chunk /= 2) {
const uint64_t need = strata::prefill::Prefill::bytes_needed(g, ss, chunk);
const int64_t blob = (int64_t) strata::kernels::cpu::expert_layout().max_blob;
int64_t k = (int64_t) ((need + (uint64_t) blob - 1) / (uint64_t) blob);
if (xcache.slot_offsets() != nullptr) {
@@ -1783,16 +1970,27 @@ int main(int argc, char** argv) {
(uint64_t) (xcache.bytes() - (int64_t) xcache.slot_offsets()[xcache.slots() - k]) < need) ++k;
}
if (k + 128 <= xcache.slots()) {
if (chunk != o.prefill_chunk)
std::fprintf(stderr, "strata serve: prompt chunk %lld -> %lld tokens so its buffers fit in the "
"expert cache\n", (long long) o.prefill_chunk, (long long) chunk);
o.prefill_chunk = chunk;
lend_first = (int32_t) (xcache.slots() - k);
borrow = xcache.device_slot(lend_first);
borrow_bytes = xcache.slot_offsets() ? (uint64_t) (xcache.bytes() - (int64_t) xcache.slot_offsets()[lend_first])
: (uint64_t) k * (uint64_t) blob;
break;
}
}
if (borrow != nullptr)
std::fprintf(stderr, "strata serve: the prompt path borrows %lld cache slots (%.2f GiB)\n",
(long long) (xcache.slots() - lend_first), (double) borrow_bytes / 1073741824.0);
else
std::fprintf(stderr, "strata serve: the prompt path allocates its own buffers (too few cache slots to borrow)\n");
if (!sp.init(wt, g, ss, srcp, &xcache, host_res.data(), o.prefill_chunk, main_cs, err, borrow, borrow_bytes)) {
std::fprintf(stderr, "strata serve: %s\n", err.c_str());
return 1;
}
mem_mark("the head and the prompt path");
strata::core::Verifier ver;
strata::core::VerifyHits vh;
vh.d_res = thits.d_res;
@@ -1803,14 +2001,52 @@ int main(int argc, char** argv) {
std::fprintf(stderr, "strata serve: %s\n", err.c_str());
return 1;
}
mem_mark("the verifier and the drafter's binding");
ver.set_split(o.spec_split);
ver.set_pcie_mode(o.pcie_mode == "dma" ? 0 : o.pcie_mode == "direct" ? 1 : o.pcie_mode == "kernel" ? 2
: native_pack ? 0 : 2); // auto: DMA for the native packs, the copy kernel for Q2_0
std::vector<int64_t> cur;
// ---- the conversation cache (see ConvCheckpoint). `live` is what the session holds right now: the tokens
// it has consumed, so a request that starts with exactly them continues without any copy. `checks` are the
// saved points; every one of them is a prefix of `live` (the loop drops the rest).
std::vector<int32_t> live;
std::vector<ImgKey> live_imgs, req_imgs;
bool live_ok = false;
std::vector<ConvCheckpoint> checks;
int64_t pp_total = 0, pp_from = 0, pp_next_check = 0;
Clock::time_point pp_t0 = Clock::now();
auto imgs_below = [&](const std::vector<ImgKey>& all, int64_t L) {
std::vector<ImgKey> v;
for (const ImgKey& k : all) if (k.start < L) v.push_back(k);
return v;
};
// a checkpoint of the state after `cur[0, L)`; false only when the copy itself failed
auto checkpoint_at = [&](int64_t L) -> bool {
if (o.prompt_cache <= 0 || L < 1) return true;
for (const ConvCheckpoint& c : checks) if ((int64_t) c.ids.size() == L) return true;
ConvCheckpoint c;
c.ids.assign(cur.begin(), cur.begin() + L);
c.imgs = imgs_below(req_imgs, L);
if (cudaDeviceSynchronize() != cudaSuccess || !checkpoint_save(c, ss, g)) return false;
checks.push_back(std::move(c));
while ((int) checks.size() > o.prompt_cache) checks.erase(checks.begin()); // the oldest goes first
return true;
};
sp.on_chunk = [&](const float* R_rows, int64_t T, int64_t p0, std::string& e) -> bool {
std::vector<int32_t> nxt((size_t) T);
for (int64_t t = 0; t < T; ++t) nxt[(size_t) t] = (int32_t) cur[(size_t) (p0 + t + 1)];
return mtp.prefill(R_rows, nxt.data(), T, p0, e);
if (!mtp.prefill(R_rows, nxt.data(), T, p0, e)) return false;
// progress for the server window: PP <position reached> <prompt tokens> <ms> <fresh tokens/s>
const int64_t done = p0 + T;
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);
std::fflush(stdout);
if (o.prompt_cache_every > 0 && done >= pp_next_check) {
if (!checkpoint_at(done)) { e = "saving a conversation checkpoint failed"; return false; }
pp_next_check = done + o.prompt_cache_every;
}
return true;
};
drive.d.plan = ver.plan_sink();
drive.d.pcie_num = std::max(0, std::min(256, (int) (o.pcie_frac * 256.0 + 0.5)));
@@ -1875,11 +2111,59 @@ int main(int argc, char** argv) {
for (float& v : drive.d.usage) v *= 0.7f;
return true;
};
std::printf("READY %lld\n", (long long) o.max_context);
// 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};
std::mutex in_mu;
std::condition_variable in_cv;
std::deque<std::string> in_lines;
bool in_eof = false;
std::thread([&] {
std::string l;
while (std::getline(std::cin, l)) {
if (!l.empty() && l.back() == '\r') l.pop_back();
if (l == "STOP") { stop_req.store(true); continue; }
std::lock_guard<std::mutex> lk(in_mu);
in_lines.push_back(l);
in_cv.notify_one();
}
std::lock_guard<std::mutex> lk(in_mu);
in_eof = true;
in_cv.notify_one();
}).detach();
auto next_line = [&](std::string& out) -> bool {
std::unique_lock<std::mutex> lk(in_mu);
in_cv.wait(lk, [&] { return !in_lines.empty() || in_eof; });
if (in_lines.empty()) return false;
out = std::move(in_lines.front());
in_lines.pop_front();
return true;
};
sp.should_stop = [&] { return stop_req.load(); };
// STRATA_TRACE=1: one stderr line per step of a request (the log shows where a request stops)
const bool trace = std::getenv("STRATA_TRACE") != nullptr;
auto tr = [&](const char* what, long long a = -1, long long b = -1) {
if (!trace) return;
std::fprintf(stderr, "strata trace: %s %lld %lld\n", what, a, b);
std::fflush(stderr);
};
{
// what is left once everything is allocated: under WDDM a GPU filled to the brim does not fail, it pages -
// and a page-in while the verify graph spins on a host flag stalls the request for good
size_t free_b = 0, total_b = 0;
cudaMemGetInfo(&free_b, &total_b);
std::fprintf(stderr, "strata serve: %lld MiB of VRAM free with everything loaded%s\n",
(long long) (free_b >> 20), free_b < ((size_t) 128 << 20)
? " - LOW: requests may stall; lower --max-context or raise --vram-reserve-mib" : "");
}
std::printf("READY %lld stop\n", (long long) o.max_context); // "stop": this engine honours STOP
std::fflush(stdout);
std::string line;
int64_t rounds = 0;
const int S = o.spec;
const int S_mtp = o.mtp_max_t > 0 ? std::min(o.mtp_max_t, S) : S; // the MTP's windows; suffixes go up to S
if (S_mtp < S) mtp.set_max_drafts(S_mtp - 1);
strata::spec::SuffixDrafter sfx(std::max(1, o.suffix_draft), 64, (size_t) o.max_context + 4096);
// The vision path (--vision): GENI <max_new> <embeddings file> <id,id,...> carries images. The file is one
// or more strata-vision records (int32 'SVE1', n, nx, ny, n_embd, then n x n_embd floats) in prompt order;
// each image's rows go to its run of <|image_pad|> tokens, whose M-RoPE positions are mtmd's: t = p,
@@ -1888,9 +2172,9 @@ int main(int argc, char** argv) {
bool mrope_identity = true;
std::vector<float> img_rows;
std::vector<const float*> row_ptr;
while (std::getline(std::cin, line)) {
if (!line.empty() && line.back() == '\r') line.pop_back();
while (next_line(line)) {
if (line == "QUIT") break;
stop_req.store(false); // a STOP that arrived between requests is stale
const bool geni = line.rfind("GENI ", 0) == 0;
if (!geni && line.rfind("GEN ", 0) != 0) {
std::printf("ERR expected: GEN <max_new> <id,id,...> or GENI <max_new> <file> <id,id,...>\n");
@@ -1911,6 +2195,7 @@ int main(int argc, char** argv) {
continue;
}
const int64_t n = (int64_t) ids.size();
req_imgs.clear();
if (geni && !o.vision) { std::printf("ERR this engine was started without --vision\n"); continue; }
if (geni || !mrope_identity) {
// positions for every cell this request can reach; the identity again for a text request
@@ -1952,6 +2237,12 @@ int main(int argc, char** argv) {
if (ids[(size_t) i] != kImagePad) { put(i, p, p, p); ++p; ++i; continue; }
if (k >= imgs.size()) { ve = "the prompt has more images than the embeddings file"; break; }
const Img& im = imgs[k++];
{ // what the conversation cache compares: a picture is its grid and its embeddings
const int64_t grid[3] = {im.n, im.nx, im.ny};
uint64_t h = fnv1a(grid, sizeof grid);
h = fnv1a(img_rows.data() + im.off, (size_t) im.n * (size_t) g.n_embd * sizeof(float), h);
req_imgs.push_back({i, h});
}
for (int64_t j = 0; j < im.n && ve.empty(); ++j)
if (i + j >= n || ids[(size_t) (i + j)] != kImagePad)
ve = "image " + std::to_string(k) + " has " + std::to_string((long long) im.n) +
@@ -1968,7 +2259,9 @@ int main(int argc, char** argv) {
if (ve.empty() && n > 0 && ids[(size_t) (n - 1)] == kImagePad) ve = "the prompt cannot end in an image";
for (int64_t c = n; ve.empty() && c < cells; ++c) put(c, p + (c - n), p + (c - n), p + (c - n));
}
tr("positions built", (long long) img_rows.size());
cudaDeviceSynchronize();
tr("device idle");
if (ve.empty() && cudaMemcpy(d_mrope, mrope_host.data(), mrope_host.size() * sizeof(int32_t),
cudaMemcpyHostToDevice) != cudaSuccess)
ve = "the image position upload failed";
@@ -1994,65 +2287,235 @@ int main(int argc, char** argv) {
if (bad) { std::printf("ERR a token id is outside the vocabulary\n"); continue; }
cur = ids;
const Clock::time_point r0 = Clock::now();
strata::core::session_zero(ss, g, nullptr, main_cs);
cudaStreamSynchronize(main_stream);
// ---- where this request starts reading: the live session, or a checkpoint, whose tokens AND pictures are
// exactly the start of this prompt - at most n - 1 of them, the last token is always the first window
auto starts_with = [&](const std::vector<int32_t>& pre, const std::vector<ImgKey>& pre_imgs) -> bool {
const int64_t L = (int64_t) pre.size();
if (L < 1 || L > n - 1) return false;
for (int64_t i = 0; i < L; ++i)
if ((int32_t) ids[(size_t) i] != pre[(size_t) i]) return false;
return imgs_below(req_imgs, L) == pre_imgs;
};
int64_t resume = 0;
bool from_live = false;
if (o.prompt_cache > 0) {
if (live_ok && starts_with(live, live_imgs)) { resume = (int64_t) live.size(); from_live = true; }
for (const ConvCheckpoint& c : checks)
if ((int64_t) c.ids.size() > resume && starts_with(c.ids, c.imgs)) {
resume = (int64_t) c.ids.size();
from_live = false;
}
}
// this request rewrites every cell from `resume` on, so a checkpoint past it (or not on this prompt's
// path) no longer has its cells; the ones kept are prefixes of both the old tokens and the new
checks.erase(std::remove_if(checks.begin(), checks.end(), [&](const ConvCheckpoint& c) {
return (int64_t) c.ids.size() > resume || !starts_with(c.ids, c.imgs);
}), checks.end());
live_ok = false; // until this request has finished, the session is in between
int64_t reread_to = -1; // STRATA_CKPT_REREAD only: read [0, reread_to) again instead of restoring
if (resume == 0) {
strata::core::session_zero(ss, g, nullptr, main_cs);
cudaStreamSynchronize(main_stream);
checks.clear();
} else if (!from_live) {
const ConvCheckpoint* c = nullptr;
for (const ConvCheckpoint& k : checks) if ((int64_t) k.ids.size() == resume) c = &k;
static const bool reread = std::getenv("STRATA_CKPT_REREAD") != nullptr;
if (reread && c != nullptr) {
// THE CHECK OF THE CHECKPOINT: instead of restoring it, read its tokens again from position 0 in
// one run (below, with the prompt path's slots lent like any read) - the same chunks the request
// that saved it read them in, when that request started at 0. With the VRAM expert set fixed
// (--adapt-swaps 0) the answer must match the restored one token for token; anything the
// checkpoint missed shows up as a difference.
strata::core::session_zero(ss, g, nullptr, main_cs);
cudaStreamSynchronize(main_stream);
reread_to = resume;
std::fprintf(stderr, "strata serve: STRATA_CKPT_REREAD: reading %lld tokens again instead of "
"restoring\n", (long long) resume);
} else if (c == nullptr || !checkpoint_restore(*c, ss, g)) {
std::printf("ERR restoring a conversation checkpoint failed\n");
return 1;
}
}
// KV streaming: the drafter's ring may hold cells past `resume` from a longer turn; the main layers'
// host copies and slots are always current (every writer writes both), so they need nothing
if (resume > 0 && reread_to <= 0) mtp.kv_restore(resume);
tr("request", n, geni ? 1 : 0);
mtp.set_prompt_len(n);
const int64_t read_from = reread_to > 0 ? 0 : resume;
pp_total = n;
pp_from = read_from;
pp_t0 = r0;
pp_next_check = reread_to > 0 ? INT64_MAX : resume + o.prompt_cache_every;
std::printf("RESUME %lld\n", (long long) resume); // before reading: this many prompt tokens are reused
std::fflush(stdout);
// A SHORT PART OF THE PROMPT - the new message of a chat that continues from a checkpoint, the assistant
// header - goes through the verify windows, S tokens at a time, as decode reads them. The batched path
// costs ~300 ms per run however few tokens it has (it streams every expert the chunk routes to that is
// not in VRAM over PCIe), and it borrows slots it must refill after (~180 ms); a window costs ~16 ms a
// token, with the misses on the CPU. Each part below is decided on its own, so a long first message is
// read batched and its header still goes through the windows. Picture rows need the batched path.
// STRATA_CKPT_REREAD compares a restored checkpoint with a batched re-read, so it keeps every read batched.
static const bool no_short = std::getenv("STRATA_CKPT_REREAD") != nullptr;
auto windows_ok = [&](int64_t a, int64_t b) -> bool {
if (no_short || b - a > o.short_read) return false;
if (sp.embd_rows != nullptr)
for (int64_t i = a; i < b; ++i)
if (sp.embd_rows[i] != nullptr) return false;
return true;
};
// tokens [a, b) through the windows: commit all of them, then give the draft layer their residuals
auto read_windows = [&](int64_t a, int64_t b, std::string& e) -> bool {
std::vector<int32_t> win((size_t) S), outw((size_t) S), nxt((size_t) S);
for (int64_t q = a; q < b;) {
if (stop_req.load()) { e = "cancelled"; return false; }
const int T = (int) std::min<int64_t>(S, b - q);
for (int t = 0; t < T; ++t) {
win[(size_t) t] = (int32_t) cur[(size_t) (q + t)];
nxt[(size_t) t] = (int32_t) cur[(size_t) (q + t + 1)];
}
drive.d.layers = 0;
drive.d.experts = 0;
drive.d.failed = false;
if (!ver.run(T, win.data(), q, &drive_pool_multi, &drive, outw.data(), e) || drive.d.failed) {
if (drive.d.failed && drive.d.fail) e = drive.d.fail;
return false;
}
if (!ver.commit(T, e) || !mtp.prefill(ver.final_R_all(), nxt.data(), T, q, e)) return false;
q += T;
}
const double ms = std::chrono::duration<double, std::milli>(Clock::now() - pp_t0).count();
std::printf("PP %lld %lld %.0f %.1f\n", (long long) b, (long long) pp_total, ms,
ms > 0.0 ? 1000.0 * (double) (b - pp_from) / ms : 0.0);
std::fflush(stdout);
return true;
};
// the batched path's slots are lent just before its first run and given back (refilled) before a window
// reads - so the windows always see the whole expert cache - or once the prompt is read
std::vector<std::pair<int32_t, int32_t>> lent_now;
apply_pending(true);
if (lend_first >= 0) {
auto lend = [&] {
if (lend_first < 0 || !lent_now.empty()) return;
for (size_t i = 0; i < host_res.size(); ++i)
if (host_res[i] >= lend_first) {
lent_now.emplace_back((int32_t) i, host_res[i]);
host_res[i] = strata::core::kNotResident;
}
cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice);
}
if (n > 1 && !sp.run(ids.data(), n - 1, 0, err)) {
std::fprintf(stderr, "strata serve: %s\n", err.c_str());
std::printf("ERR %s\n", err.c_str());
return 1;
}
for (const auto& [i, slot] : lent_now) {
const uint8_t* b = srcp->blob(i / g.n_expert, i % g.n_expert);
if (b == nullptr || !xcache.fill_slot_blocking(slot, b, err,
(int64_t) strata::kernels::cpu::expert_layout().blob_bytes(i / g.n_expert))) {
};
auto refill = [&](std::string& e) -> bool {
if (lent_now.empty()) return true;
tr("refill start", (long long) lent_now.size());
for (const auto& [i, slot] : lent_now) {
const uint8_t* b = srcp->blob(i / g.n_expert, i % g.n_expert);
if (b == nullptr || !xcache.fill_slot_blocking(slot, b, e,
(int64_t) strata::kernels::cpu::expert_layout().blob_bytes(i / g.n_expert)))
return false;
host_res[(size_t) i] = slot;
}
cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice);
lent_now.clear();
return true;
};
apply_pending(true);
bool cancelled = false;
tr("prompt start", n - 1);
// The prompt is read in two parts when it has a turn boundary past `resume`: up to the last <|im_start|>
// (the conversation so far), a checkpoint there, then the new turn's header. The next request of the same
// chat renders the same history - but not always the same header or the thinking of this reply - so that
// checkpoint is the one it reuses.
int64_t turn_at = -1;
if (o.prompt_cache > 0 && o.turn_token >= 0)
for (int64_t i = n - 1; i > resume; --i)
if (ids[(size_t) i] == o.turn_token) { turn_at = i; break; }
int64_t at = read_from;
for (const int64_t to : {reread_to, turn_at, n - 1}) {
if (to <= at) continue;
const bool win = windows_ok(at, to);
if (win && !refill(err)) {
std::printf("ERR refilling a lent slot failed: %s\n", err.c_str());
return 1;
}
host_res[(size_t) i] = slot;
if (!win) lend();
const auto tsp = Clock::now();
const bool sp_ok = win ? read_windows(at, to, err) : sp.run(ids.data() + at, to - at, at, err);
if (trace) {
std::fprintf(stderr, "strata trace: read %lld tokens (%s) in %.1f ms\n", (long long) (to - at),
win ? "windows" : "batched",
std::chrono::duration<double, std::milli>(Clock::now() - tsp).count());
std::fflush(stderr);
}
if (!sp_ok) {
if (!stop_req.load()) {
std::fprintf(stderr, "strata serve: %s\n", err.c_str());
std::printf("ERR %s\n", err.c_str());
return 1;
}
cancelled = true; // stopped while reading the prompt: refill the lent slots below, then DONE cancel
break;
}
at = to;
if (to == turn_at && !checkpoint_at(to)) {
std::printf("ERR saving a conversation checkpoint failed\n");
return 1;
}
}
if (!lent_now.empty())
cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice);
if (!refill(err)) {
std::printf("ERR refilling a lent slot failed: %s\n", err.c_str());
return 1;
}
tr("prompt done (slots refilled)");
const double prompt_ms = std::chrono::duration<double, std::milli>(Clock::now() - r0).count();
std::printf("REUSED %lld\n", (long long) resume); // the prompt is read; the first window comes next
std::fflush(stdout);
// the verify windows: the first holds the last prompt token alone
int64_t p = n - 1;
int32_t x = (int32_t) ids[(size_t) (n - 1)];
std::vector<int32_t> drafts((size_t) S, 0), window((size_t) S), outv((size_t) S);
std::vector<float> dprob((size_t) S, 0.0f);
std::vector<int32_t> sbuf((size_t) S, 0);
if (o.suffix_draft > 0) {
sfx.reset();
for (int64_t t : ids) sfx.append((int32_t) t);
}
bool first_window = true;
int64_t produced_n = 0;
int64_t produced_n = 0, sfx_windows = 0, sfx_drafts = 0, sfx_ok = 0;
int64_t draft_offered = 0, draft_accepted = 0;
// what the session holds once this request is done: the prompt read so far, then every committed token
std::vector<int32_t> consumed;
consumed.reserve((size_t) (n + max_new + S));
for (int64_t i = 0; i < n - 1; ++i) consumed.push_back((int32_t) ids[(size_t) i]);
const char* finish = "length";
const Clock::time_point d0 = Clock::now();
while (produced_n < max_new) {
int T = S;
if (cancelled) finish = "cancel";
while (!cancelled && produced_n < max_new) {
int T = S_mtp;
if (o.spec_min_p > 0.0) {
T = 1;
while (T < S && dprob[(size_t) T - 1] >= (float) o.spec_min_p) ++T;
while (T < S_mtp && dprob[(size_t) T - 1] >= (float) o.spec_min_p) ++T;
}
if (first_window) T = 1;
// a repeat of earlier context proposes more than the MTP is confident about: take it, but only where the
// MTP's own first guess agrees (a repeat the model is not following costs a whole window)
bool from_sfx = false;
if (o.suffix_draft > 0 && !first_window) {
const int k = sfx.propose(S - 1, sbuf.data());
if (k + 1 > T && sbuf[0] == drafts[0]) { T = k + 1; from_sfx = true; }
}
if (p + T > o.max_context) break;
window[0] = x;
for (int i = 1; i < T; ++i) window[(size_t) i] = drafts[(size_t) i - 1];
for (int i = 1; i < T; ++i) window[(size_t) i] = from_sfx ? sbuf[(size_t) i - 1] : drafts[(size_t) i - 1];
drive.d.layers = 0;
drive.d.experts = 0;
drive.d.failed = false;
apply_pending(false);
tr("window", p, T);
if (!ver.run(T, window.data(), p, &drive_pool_multi, &drive, outv.data(), err) || drive.d.failed) {
std::printf("ERR %s\n", drive.d.failed && drive.d.fail ? drive.d.fail : err.c_str());
return 1;
}
int a = 0;
while (a < T - 1 && window[(size_t) a + 1] == outv[(size_t) a]) ++a;
if (from_sfx) { ++sfx_windows; sfx_drafts += T - 1; sfx_ok += a; }
std::thread adapt_thr; // the adaptive tier beside the commit and the draft (as in generate)
bool adapt_ok = true;
if (!drive.d.usage.empty() && ((rounds + 1) % o.adapt_every) == 0)
@@ -2062,11 +2525,16 @@ int main(int argc, char** argv) {
std::printf("ERR %s\n", err.c_str());
return 1;
}
// the window's first a + 1 tokens are in the session now (the last output is not: it is next x)
for (int i = 0; i <= a; ++i) consumed.push_back(window[(size_t) i]);
draft_offered += T - 1;
draft_accepted += a;
first_window = false;
bool eos = false;
for (int i = 0; i <= a && produced_n < max_new && !eos; ++i) {
std::printf("T %d\n", (int) outv[(size_t) i]);
++produced_n;
if (o.suffix_draft > 0) sfx.append(outv[(size_t) i]);
eos = std::find(o.eos_ids.begin(), o.eos_ids.end(), (int64_t) outv[(size_t) i]) != o.eos_ids.end();
}
std::fflush(stdout);
@@ -2083,16 +2551,118 @@ int main(int argc, char** argv) {
return 1;
}
if (eos) { finish = "stop"; break; }
if (stop_req.load()) { finish = "cancel"; break; }
x = outv[(size_t) a];
p += a + 1;
}
const double decode_ms = std::chrono::duration<double, std::milli>(Clock::now() - d0).count();
std::printf("DONE %lld %lld %.1f %.1f %s\n", (long long) produced_n, (long long) n, prompt_ms, decode_ms,
finish);
if (!cancelled) {
// a prompt stopped halfway leaves the session somewhere between two chunks: nothing to continue from
// (the checkpoints taken while reading it are still good)
live.swap(consumed);
live_imgs = imgs_below(req_imgs, (int64_t) live.size());
live_ok = o.prompt_cache > 0;
}
static const bool state_hash = std::getenv("STRATA_STATE_HASH") != nullptr;
if (state_hash && live_ok) {
// DEBUG: a fingerprint of every part of the session over the positions it holds ([0, L)), and
// separately of what lies past them in the last KV page (stale cells, fine unless something reads them)
cudaDeviceSynchronize();
const int64_t L = (int64_t) live.size();
const strata::kernels::QsaShapes qs = [&] {
strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes();
s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_dim = g.idx_key_dim;
return s;
}();
auto hash_dev = [&](const void* p, size_t bytes, uint64_t h) {
std::vector<uint8_t> b(bytes);
if (bytes) cudaMemcpy(b.data(), p, bytes, cudaMemcpyDefault); // VRAM or a streamed host copy
return fnv1a(b.data(), b.size(), h);
};
// the cells [c0, c1) of one int8 K or V pool ([page][kv_head][page_size][head_dim]), bytes per value `w`
auto hash_cells = [&](const void* pool, int64_t per_cell, int64_t c0, int64_t c1, uint64_t h) {
const int64_t ps = qs.page_size;
for (int64_t pg = c0 / ps; pg * ps < c1; ++pg)
for (int64_t hd = 0; hd < qs.n_head_kv; ++hd) {
const int64_t a = std::max(c0, pg * ps) - pg * ps, e = std::min(c1, (pg + 1) * ps) - pg * ps;
const size_t off = (size_t) (((pg * qs.n_head_kv + hd) * ps + a) * per_cell);
h = hash_dev((const uint8_t*) pool + off, (size_t) ((e - a) * per_cell), h);
}
return h;
};
const ConvStateSizes z = conv_state_sizes(g);
uint64_t h_gdn = hash_dev(ss.gdn_state, z.gdn, 1469598103934665603ull);
if (std::getenv("STRATA_STATE_HASH_GDN") != nullptr) { // per GDN layer: which one differs first
const size_t per = z.gdn / (size_t) g.n_gdn_layers();
std::string s;
char b[8];
for (int64_t i = 0; i < g.n_gdn_layers(); ++i) {
std::snprintf(b, sizeof(b), "%04llx ", (unsigned long long) (hash_dev((const uint8_t*) ss.gdn_state + i * per, per, 1469598103934665603ull) & 0xffff));
s += b;
}
std::fprintf(stderr, "strata serve: STATE_HASH_GDN %s\n", s.c_str());
}
uint64_t h_ple = hash_dev(ss.ple_hist, z.ple, 1469598103934665603ull);
uint64_t h_tail = 1469598103934665603ull, h_pool = h_tail, h_kv = h_tail, h_stale = h_tail;
const int64_t kvb = qs.head_dim, scb = (qs.head_dim / 64) * 2;
const int64_t end_cell = std::min<int64_t>(((L + qs.page_size - 1) / qs.page_size) * qs.page_size,
ss.qsa_states[0].max_cells);
for (int64_t i = 0; i < g.n_qsa_layers(); ++i) {
const strata::core::QsaState& st = ss.qsa_states[i];
h_tail = hash_dev(st.idx_tail, z.tail, h_tail);
h_pool = hash_dev(st.idx_pooled, (size_t) (L / qs.idx_block) * qs.idx_dim * 4, h_pool);
// KV streaming: the host copy is the identity layout and holds every cell
const bool hs = st.kv_mode != 0;
for (const auto& [pool, w] : {std::pair<const void*, int64_t>{hs ? st.host.k_q : st.k_q, kvb},
{hs ? st.host.v_q : st.v_q, kvb},
{hs ? st.host.k_scale : st.k_scale, scb},
{hs ? st.host.v_scale : st.v_scale, scb}}) {
h_kv = hash_cells(pool, w, 0, L, h_kv);
h_stale = hash_cells(pool, w, L, end_cell, h_stale);
}
}
const strata::core::QsaState& ms = mtp.kv_state();
uint64_t h_mtp = 1469598103934665603ull;
const int64_t mL = std::min<int64_t>(L, ms.max_cells);
const bool mh = ms.kv_mode != 0;
for (const auto& [pool, w] : {std::pair<const void*, int64_t>{mh ? ms.host.k_q : ms.k_q, kvb},
{mh ? ms.host.v_q : ms.v_q, kvb},
{mh ? ms.host.k_scale : ms.k_scale, scb},
{mh ? ms.host.v_scale : ms.v_scale, scb}})
if (pool != nullptr) h_mtp = hash_cells(pool, w, 0, mL, h_mtp);
std::fprintf(stderr, "strata serve: STATE_HASH L=%lld gdn=%016llx ple=%016llx tail=%016llx pooled=%016llx "
"kv=%016llx mtp=%016llx stale=%016llx ple_prev=%d,%d\n", (long long) L,
(unsigned long long) h_gdn, (unsigned long long) h_ple, (unsigned long long) h_tail,
(unsigned long long) h_pool, (unsigned long long) h_kv, (unsigned long long) h_mtp,
(unsigned long long) h_stale, ss.ple_prev[0], ss.ple_prev[1]);
}
// DONE <generated> <prompt> <prompt ms> <decode ms> <finish> <drafts accepted> <drafts offered> <reused>
std::printf("DONE %lld %lld %.1f %.1f %s %lld %lld %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);
std::fflush(stdout);
std::fprintf(stderr, "strata serve: %lld prompt tokens in %.0f ms (%.1f tok/s), %lld generated in %.0f ms "
"(%.1f tok/s)\n", (long long) n, prompt_ms, prompt_ms > 0 ? 1000.0 * n / prompt_ms : 0.0,
(long long) produced_n, decode_ms, decode_ms > 0 ? 1000.0 * produced_n / decode_ms : 0.0);
const int64_t fresh = n - resume;
std::fprintf(stderr, "strata serve: prompt %lld tokens = %lld reused + %lld 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,
decode_ms > 0 ? 1000.0 * produced_n / decode_ms : 0.0, (long long) draft_accepted,
(long long) draft_offered, checks.size(), cancelled ? " (cancelled)" : "");
if (g.n_qsa_layers() > 0 && ss.qsa_states[0].kv_mode == 1) {
// KV streaming, cumulative over the process: blocks the selections named vs blocks read from RAM
uint64_t miss = 0, look = 0;
bool over = false;
for (int64_t i = 0; i < g.n_qsa_layers(); ++i) {
const strata::kernels::KvStreamCounters c = strata::kernels::kv_stream_counters(ss.qsa_states[i].map);
miss += c.misses; look += c.lookups; over = over || c.overflow;
}
std::fprintf(stderr, "strata serve: KV streaming: %.2f%% of %llu block reads hit VRAM, %.1f MiB read "
"from RAM%s\n", look ? 100.0 * (double) (look - miss) / (double) look : 100.0,
(unsigned long long) look, (double) miss * 4224.0 / 1048576.0,
over ? " - OVERFLOW (too few resident cells)" : "");
}
if (sfx_windows > 0)
std::fprintf(stderr, "strata serve: suffix drafts: %lld windows, %lld of %lld drafts accepted\n",
(long long) sfx_windows, (long long) sfx_ok, (long long) sfx_drafts);
}
return 0;
}
@@ -2426,6 +2996,7 @@ int main(int argc, char** argv) {
"--expert-cache and the token graph)\n");
return 2;
}
mem_mark("the head and the prompt path");
strata::core::Verifier ver;
strata::core::VerifyHits vh;
vh.d_res = thits.d_res;
@@ -2440,6 +3011,7 @@ int main(int argc, char** argv) {
std::fprintf(stderr, "strata generate: %s\n", err.c_str());
return 1;
}
mem_mark("the verifier and the drafter's binding");
ver.set_split(o.spec_split);
ver.set_pcie_mode(o.pcie_mode == "dma" ? 0 : o.pcie_mode == "direct" ? 1 : o.pcie_mode == "kernel" ? 2
: native_pack ? 0 : 2); // auto: DMA for the native packs, the copy kernel for Q2_0
@@ -2535,16 +3107,30 @@ int main(int argc, char** argv) {
std::vector<int32_t> window((size_t) o.spec), outv((size_t) o.spec);
std::vector<int64_t> accepted_hist((size_t) o.spec, 0);
int64_t rounds = 0, drafts_total = 0, drafts_ok = 0, corrupt_counter = 0;
const int S_mtp = o.mtp_max_t > 0 ? std::min(o.mtp_max_t, o.spec) : o.spec;
if (use_mtp && S_mtp < o.spec) mtp.set_max_drafts(S_mtp - 1);
strata::spec::SuffixDrafter sfx(std::max(1, o.suffix_draft), 64, (size_t) o.max_context + 4096);
std::vector<int32_t> sbuf((size_t) o.spec, 0);
int64_t sfx_windows = 0, sfx_drafts = 0, sfx_ok = 0;
if (o.suffix_draft > 0) {
for (int64_t t : o.tokens) sfx.append((int32_t) t);
for (int64_t t : produced) sfx.append((int32_t) t);
}
const double pool_ms0 = drive.cpu_ms;
const int64_t misses0 = drive.d.multi_misses, entries0 = drive.d.multi_entries;
while ((int64_t) produced.size() < o.max_new) {
const Clock::time_point t0 = Clock::now();
int T = o.spec;
int T = S_mtp;
if (use_mtp && o.spec_min_p > 0.0) {
T = 1;
while (T < o.spec && dprob[(size_t) T - 1] >= (float) o.spec_min_p) ++T;
while (T < S_mtp && dprob[(size_t) T - 1] >= (float) o.spec_min_p) ++T;
}
if (first_window) T = 1;
bool from_sfx = false;
if (o.suffix_draft > 0 && !first_window) {
const int k = sfx.propose(o.spec - 1, sbuf.data());
if (k + 1 > T && (!use_mtp || sbuf[0] == drafts[0])) { T = k + 1; from_sfx = true; }
}
++window_hist[(size_t) T];
if (p + T > o.max_context) {
std::fprintf(stderr, "strata generate: ran out of context at position %lld\n", (long long) p);
@@ -2553,7 +3139,8 @@ int main(int argc, char** argv) {
window[0] = x;
for (int i = 1; i < T; ++i) {
const size_t at = produced.size() - 1 + (size_t) i;
int32_t d = use_mtp ? drafts[(size_t) i - 1] : at < oracle.size() ? (int32_t) oracle[at] : 0;
int32_t d = from_sfx ? sbuf[(size_t) i - 1] : use_mtp ? drafts[(size_t) i - 1]
: at < oracle.size() ? (int32_t) oracle[at] : 0;
if (o.spec_corrupt > 0 && (++corrupt_counter % o.spec_corrupt) == 0) d = (d + 1) % (int32_t) n_vocab;
window[(size_t) i] = d;
}
@@ -2592,9 +3179,11 @@ int main(int argc, char** argv) {
drafts_total += T - 1;
drafts_ok += a;
++accepted_hist[(size_t) a];
if (from_sfx) { ++sfx_windows; sfx_drafts += T - 1; sfx_ok += a; }
bool eos = false;
for (int i = 0; i <= a && (int64_t) produced.size() < o.max_new && !eos; ++i) {
produced.push_back(outv[(size_t) i]);
if (o.suffix_draft > 0) sfx.append(outv[(size_t) i]);
eos = o.stop_eos && std::find(o.eos_ids.begin(), o.eos_ids.end(), (int64_t) outv[(size_t) i]) != o.eos_ids.end();
}
if (eos) {
@@ -2626,6 +3215,9 @@ int main(int argc, char** argv) {
for (size_t i = 1; i < window_hist.size(); ++i) std::printf(" T%zu:%lld", i, (long long) window_hist[i]);
std::printf(" (min draft probability %.2f)\n", o.spec_min_p);
}
if (o.suffix_draft > 0)
std::printf("%-24s %lld windows, drafts accepted %lld of %lld\n", "suffix drafts", (long long) sfx_windows,
(long long) sfx_ok, (long long) sfx_drafts);
std::printf("%-24s", "accepted per round");
for (size_t i = 0; i < accepted_hist.size(); ++i) std::printf(" %zu:%lld", i, (long long) accepted_hist[i]);
std::printf("\n");
+24
View File
@@ -99,6 +99,30 @@ int main(int argc, char** argv) {
if (gg) gguf_free(gg);
}
if (n_embd <= 0) { std::printf("ERR the vision encoder has no projection_dim\n"); std::fflush(stdout); return 1; }
// Warm up at the LARGEST picture before READY: the encoder's GPU work buffers are allocated now, not at the first
// real picture. The server starts this process before the engine, so the engine sizes its expert cache from
// what is really left; allocating ~1 GB later, on a GPU the engine has filled, made Windows page GPU memory and
// the engine crawl to a standstill. (A square image well above any cap; mtmd scales it to the token limit.)
{
const uint32_t side = 2048;
std::vector<unsigned char> rgb((size_t) side * side * 3, 128);
mtmd_bitmap* bm = mtmd_bitmap_init(side, side, rgb.data());
mtmd_input_chunks* chunks = mtmd_input_chunks_init();
const std::string marker = mtmd_default_marker();
mtmd_input_text txt{marker.c_str(), marker.size(), false, true};
const mtmd_bitmap* bms[1] = {bm};
int warm_tokens = 0;
if (bm && mtmd_tokenize(ctx, chunks, &txt, bms, 1) == 0) {
for (size_t c = 0; c < mtmd_input_chunks_size(chunks); ++c) {
const mtmd_input_chunk* ch = mtmd_input_chunks_get(chunks, c);
if (mtmd_input_chunk_get_type(ch) == MTMD_INPUT_CHUNK_TYPE_IMAGE && mtmd_encode_chunk(ctx, ch) == 0)
warm_tokens = (int) mtmd_input_chunk_get_n_tokens(ch);
}
}
std::fprintf(stderr, "strata-vision: warmed up at %d image tokens\n", warm_tokens);
mtmd_input_chunks_free(chunks);
if (bm) mtmd_bitmap_free(bm);
}
std::printf("READY %d\n", n_embd);
std::fflush(stdout);