Compare commits

...
6 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
22 changed files with 1305 additions and 116 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)
+4 -2
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:**
+12 -3
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.
@@ -310,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);
+6
View File
@@ -45,9 +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);
@@ -84,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;
+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 -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) {
+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;
}
+110 -12
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.
@@ -248,6 +250,12 @@ struct Options {
/// 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() {
@@ -267,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"
@@ -311,6 +322,8 @@ void usage() {
" --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"
@@ -603,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;
@@ -658,6 +672,8 @@ int main(int argc, char** argv) {
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;
@@ -711,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()) {
@@ -941,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
@@ -2136,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,
@@ -2309,6 +2337,9 @@ int main(int argc, char** argv) {
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;
@@ -2441,8 +2472,13 @@ int main(int argc, char** argv) {
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;
@@ -2452,15 +2488,22 @@ int main(int argc, char** argv) {
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;
@@ -2472,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)
@@ -2490,6 +2534,7 @@ int main(int argc, char** argv) {
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);
@@ -2531,7 +2576,7 @@ int main(int argc, char** argv) {
}();
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, cudaMemcpyDeviceToHost);
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`
@@ -2547,6 +2592,16 @@ int main(int argc, char** argv) {
};
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;
@@ -2556,8 +2611,12 @@ int main(int argc, char** argv) {
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);
for (const auto& [pool, w] : {std::pair<const void*, int64_t>{st.k_q, kvb}, {st.v_q, kvb},
{st.k_scale, scb}, {st.v_scale, scb}}) {
// 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);
}
@@ -2565,8 +2624,11 @@ int main(int argc, char** argv) {
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);
for (const auto& [pool, w] : {std::pair<const void*, int64_t>{ms.k_q, kvb}, {ms.v_q, kvb},
{ms.k_scale, scb}, {ms.v_scale, scb}})
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,
@@ -2585,6 +2647,22 @@ int main(int argc, char** argv) {
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;
}
@@ -3029,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);
@@ -3047,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;
}
@@ -3086,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) {
@@ -3120,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");