Compare commits

...
15 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
25 changed files with 1817 additions and 169 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)
+6 -4
View File
@@ -39,7 +39,7 @@ model - is set up for you.
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 or IQ3_XXS - see [which model](#which-model-should-i-pick)
- **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
@@ -60,6 +60,7 @@ Measured on an RTX 5070 (12 GB), a Ryzen 5 7600 and 64 GB of RAM:
| **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 |
- **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).
@@ -76,7 +77,8 @@ A card with more VRAM is faster, because more of the model fits on the GPU: an R
| --- | ---: | --- | --- | --- |
| **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 | best | you want the best answers (uses 43 GB of your 64 GB RAM) |
| **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 version:**
@@ -98,8 +100,8 @@ Not sure? Take **IQ2_XS**. You can add another one later with `START-HERE.bat --
- **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).
**Good to know:** it answers one request at a time, and it re-reads the whole conversation for every answer. So in
very long chats you wait longer before it starts writing: about 1 minute per 30,000 tokens of conversation.
**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?
+20 -6
View File
@@ -21,6 +21,7 @@ tokens, MTP speculative decoding on. "262K" is the model's full context window (
| **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)
@@ -29,9 +30,17 @@ tokens, MTP speculative decoding on. "262K" is the model's full context window (
| **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 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.
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.
@@ -200,9 +209,14 @@ print(r.choices[0].message.content)
add `"api_key": "some-long-secret"` to `strata-<model>.json` (or set the `STRATA_API_KEY` environment variable);
clients then send it as their API key.
**Current limits (v1):** one request at a time; greedy decoding (temperature is ignored); every request processes its
whole prompt again (no conversation cache yet, so long chats have a long time-to-first-token); images only when set up
with them (below); no video.
**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.
---
@@ -305,7 +319,7 @@ test images.
<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, and an **expert cache** that fills the rest of VRAM with the most-used experts (it adapts to
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).
+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);
+60 -18
View File
@@ -127,6 +127,10 @@ class StrataEngine:
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"):
self._parse_done(line)
done = True
@@ -370,7 +374,7 @@ class Service:
s["phase"] = "tool call complete"
s["tail"] = (s["tail"] + (ev.text or ""))[-600:]
def _progress(self, last_print, every=15.0):
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:
@@ -390,6 +394,7 @@ class Service:
"""Yields ("event", Event) as text arrives, then ("done", {"finish": .., "completion_tokens": ..})."""
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
@@ -402,35 +407,66 @@ class Service:
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:
if t is None: # heartbeat while the engine is quiet
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)
yield "ping", None
continue
n += 1
if t in self.stop_ids:
finish = "stop"
break
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"
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"):
el = time.time() - self.status.get("started", time.time())
print(f"[strata] done: {n} tokens in {el:.0f} s ({finish})", flush=True)
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())
@@ -668,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"):
@@ -689,7 +728,10 @@ def make_handler(svc: Service):
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"):
+60 -8
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, 2) # split models (Swift 1.5), STOP, cache sized after the slots are written
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.
@@ -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()
@@ -627,6 +667,8 @@ 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()
@@ -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:
@@ -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"),
+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);
+140 -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
}
@@ -369,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);
@@ -395,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,
@@ -533,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;
}
+501 -37
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>
@@ -116,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.
@@ -235,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() {
@@ -254,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"
@@ -292,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"
@@ -442,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) {
@@ -506,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;
@@ -557,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;
@@ -610,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()) {
@@ -840,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
@@ -1812,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 "
@@ -1876,10 +2006,47 @@ int main(int argc, char** argv) {
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)));
@@ -1994,6 +2161,9 @@ int main(int argc, char** argv) {
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,
@@ -2025,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
@@ -2066,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) +
@@ -2110,63 +2287,223 @@ 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);
}
};
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);
if (n > 1 && !sp.run(ids.data(), n - 1, 0, err)) {
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
}
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))) {
// 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 (!refill(err)) {
std::printf("ERR refilling a lent slot failed: %s\n", err.c_str());
return 1;
}
if (!lent_now.empty())
cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice);
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();
if (cancelled) finish = "cancel";
while (!cancelled && produced_n < max_new) {
int T = S;
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;
@@ -2178,6 +2515,7 @@ int main(int argc, char** argv) {
}
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)
@@ -2187,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);
@@ -2213,12 +2556,113 @@ int main(int argc, char** argv) {
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;
}
@@ -2663,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);
@@ -2681,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;
}
@@ -2720,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) {
@@ -2754,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");