Winery-Strata / src /core /verify.cpp
Zandy-Wandy's picture
🍷 Winery Strata: fork of Niko1221/Strata (Winery Cuvée family, wine-red UI)
bbb6388 verified
Raw History Blame Contribute Delete
77.2 kB
// src/core/verify.cpp - see include/strata/core/verify.hpp.
#include "strata/core/verify.hpp"
#if defined(_WIN32)
#include <intrin.h>
#endif
#include "strata/core/native_head.hpp"
#include "strata/core/on_device.hpp"
#include "strata/core/peer_experts.hpp"
#include "strata/kernels/iq_kernels.hpp"
#include "strata/kernels/cpu/expert_layout.hpp"
#include "strata/kernels/bf16_gemv.hpp"
#include "strata/kernels/native_router.hpp"
#include "strata/kernels/native_moe.hpp"
#include "strata/kernels/cpu/expert.hpp"
#include "strata/kernels/elementwise.hpp"
#include "strata/kernels/fused_gr.hpp"
#include "strata/kernels/cvec.hpp"
#include "strata/kernels/gr.hpp"
#include "strata/kernels/kv_q4.hpp"
#include "strata/kernels/kv_q8.hpp"
#include "strata/kernels/native_mmvq.hpp"
#include "strata/kernels/native_qsa.hpp"
#include "strata/kernels/native_qsa_indexer.hpp"
#include "strata/kernels/native_rope.hpp"
#include "strata/kernels/ngram.hpp"
#include "strata/kernels/ple.hpp"
#include "strata/kernels/qsa.hpp"
#include "strata/kernels/qsa_decode_attn.hpp"
#include "strata/kernels/qsa_select.hpp"
#include "strata/kernels/quantize_act.hpp"
#include "strata/kernels/rope.hpp"
#include "strata/kernels/s2_expert_grouped.hpp"
#include "strata/kernels/sampler.hpp"
#include "strata/core/progress.hpp"
#include "strata/kernels/shared_expert.hpp"
#include "strata/kernels/verify_kernels.hpp"
#include <algorithm>
#include <atomic>
#include <map>
#include <string>
#include <thread>
#include <vector>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <exception>
#include <immintrin.h>
namespace strata::core {
namespace {
constexpr float EPS = 1e-6f;
using Clock = std::chrono::steady_clock;
double ms_since(Clock::time_point t) { return std::chrono::duration<double, std::milli>(Clock::now() - t).count(); }
const bool g_dbg = std::getenv("STRATA_VERIFY_DEBUG") != nullptr;
#define VDBG(...) do { if (g_dbg) { std::fprintf(stderr, "verify dbg: " __VA_ARGS__); std::fflush(stderr); } } while (0)
struct Bump {
uint8_t* base = nullptr;
uint64_t used = 0;
template <typename T> T* take(uint64_t n) {
T* p = base ? (T*) (base + used) : nullptr;
used += (n * sizeof(T) + 255) & ~255ull;
return p;
}
};
bool mapped(size_t bytes, void** h, void** d) {
if (cudaHostAlloc(h, bytes, cudaHostAllocMapped | (peer_portable() ? cudaHostAllocPortable : 0)) != cudaSuccess) return false; // multi-GPU: portable only with a peer - the peer card writes its rows into them
std::memset(*h, 0, bytes);
return cudaHostGetDevicePointer(d, *h, 0) == cudaSuccess;
}
strata::kernels::QsaShapes shapes_of(const ModelGeometry& g) {
strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes();
s.n_head = g.n_head;
s.n_head_kv = g.n_head_kv;
s.head_dim = g.head_dim;
s.idx_n_head = g.idx_q_heads;
s.idx_dim = g.idx_key_dim;
return s;
}
const WeightRef* need(const LayerView& v, const char* suffix, std::string& err) {
const WeightRef* r = v.get(suffix);
if (r == nullptr && err.empty()) err = v.name(suffix) + " is missing";
return r;
}
bool native_of(const WeightRef* w, const std::string& name, std::string& err) {
if (w == nullptr) return false;
if (w->native_data == nullptr) {
err = "verify: " + name + " is not served natively (run with --native)";
return false;
}
return true;
}
} // namespace
namespace {
std::atomic<const Verifier*> g_diag_verifier{nullptr};
void diag_active_verifier(std::FILE* f) {
if (const Verifier* v = g_diag_verifier.load()) v->diag(f);
}
// #267: every live verifier (a layer split has one per stage), for the release before the engine ends
constexpr int kLiveMax = 16;
std::atomic<Verifier*> g_live[kLiveMax];
void release_live_verifiers(std::FILE* f) {
for (auto& slot : g_live)
if (Verifier* v = slot.load()) {
const Clock::time_point t0 = Clock::now();
const bool done = v->release_gpu_waits(5000);
if (f != nullptr)
std::fprintf(f, "strata: released the verify window's GPU waits (#267): the GPU %s\n",
done ? ("finished in " + std::to_string((long long) ms_since(t0)) + " ms").c_str()
: "did not finish within 5 s");
}
if (f != nullptr) std::fflush(f);
}
std::string released_note(bool drained) {
return drained ? "; its GPU waits were released and the GPU finished (#267)"
: "; its GPU waits were released but the GPU did not finish within 5 s (#267)";
}
// #267 test hook: STRATA_TEST_VERIFY_STALL=N withholds the last layer's flag in the N-th window (1-based), so the
// GPU spins on a flag nobody raises - the bounded window wait and the release are then what ends it. Unset: never.
const int64_t g_test_stall = [] {
const char* e = std::getenv("STRATA_TEST_VERIFY_STALL");
return e != nullptr ? (int64_t) std::atoll(e) : (int64_t) 0;
}();
} // namespace
bool Verifier::release_gpu_waits(int timeout_ms) {
released_.store(true);
// the words the spin kernels read (wait_flag_ge, wait_flag_ge_or) are mapped host memory, so a store here
// reaches them with no API call; UINT32_MAX is past every ring. (E-6's skip words are device memory, but
// wait_flag_ge_or also returns on its flag.) A host function raising flag B later only raises.
for (uint32_t* p : {h_flag_, h_flagA_, h_flagB_})
if (p != nullptr) *(volatile uint32_t*) p = UINT32_MAX;
std::atomic_thread_fence(std::memory_order_seq_cst);
_mm_sfence();
const OnDevice on_device(device_);
const Clock::time_point t0 = Clock::now();
for (cudaStream_t s : {cs_, copy_}) {
if (s == nullptr) continue;
while (cudaStreamQuery(s) == cudaErrorNotReady) {
if (ms_since(t0) > timeout_ms) return false;
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
}
return true;
}
void Verifier::diag(std::FILE* f) const {
auto rd = [](const uint32_t* p) { return p ? *(const volatile uint32_t*) p : 0u; };
// #251: outside a verify stage these are the LAST window's numbers (it finished), not the stalled work's
const char* where = progress().where.load();
const bool current = where != nullptr && std::strncmp(where, "verify window", 13) == 0;
std::fprintf(f, " verify window%s: %d tokens at position %lld, host at layer step %u; the GPU rang %u; flags: "
"served %u, plan (A) %u, copies (B) %u\n", current ? "" : " (last window, not the current stage)",
last_t_, (long long) last_pos0_, cur_layer_ + 1, rd(h_seq_), rd(h_flag_), rd(h_flagA_), rd(h_flagB_));
}
Verifier::~Verifier() {
const Verifier* self = this;
g_diag_verifier.compare_exchange_strong(self, nullptr);
for (auto& slot : g_live) {
Verifier* me = this;
slot.compare_exchange_strong(me, nullptr);
}
if (cs_) cudaStreamSynchronize(cs_);
for (auto& e : exec_)
if (e) cudaGraphExecDestroy(e);
if (commit_exec_) cudaGraphExecDestroy(commit_exec_);
if (cs_) cudaStreamDestroy(cs_);
if (copy_) { cudaStreamSynchronize(copy_); cudaStreamDestroy(copy_); }
if (commit_done_) cudaEventDestroy(commit_done_);
if (arena_) cudaFree(arena_);
void* hosts[] = {h_tok_, h_step_, h_pos_, h_commit_, h_ple_, h_out_, h_x_, h_ids_, h_w_, h_seq_, h_flag_, h_ymiss_,
h_flagA_, h_plan_, h_flagB_};
for (void* h : hosts)
if (h) cudaFreeHost(h);
}
bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& ss, const VerifyHits& hits,
const NativeHead* head, int max_t, std::string& err) {
g_diag_verifier.store(this);
diag_verify_fn().store(&diag_active_verifier);
for (auto& slot : g_live) {
Verifier* none = nullptr;
if (slot.load() == this || slot.compare_exchange_strong(none, this)) break;
}
release_gpu_fn().store(&release_live_verifiers);
cudaGetDevice(&device_); // a layer split's stage on another GPU: its streams, graphs and buffers live there
strata::kernels::fused_gr_check(); // once per card: which bitwise-equal hyper-connection read runs there
wt_ = &wt;
g_ = &g;
ss_ = &ss;
hits_ = hits;
head_ = head;
max_t_ = max_t;
sampling_.greedy = true; // a fresh verifier samples greedily until set_sampling says otherwise
sampling_.temperature = 0.0f;
if (max_t < 2 || max_t > strata::kernels::kVerifyMaxT || max_t > strata::kernels::cpu::MAXT) {
err = "verify: the window must hold 2.." + std::to_string(strata::kernels::kVerifyMaxT) + " tokens";
return false;
}
if (hits.d_res == nullptr || hits.cache_base == nullptr || hits.blob <= 0) {
err = "verify: needs the profile-filled VRAM expert tier (--expert-profile and --expert-cache); with "
"--expert-cache auto, no VRAM was left for it - the 'no VRAM is left for the expert cache' line above "
"says how much is short and what makes room (a smaller --max-context, --kv q4_0, setup --draft-vocab en, "
"images on the CPU)";
return false;
}
std::string why;
if (!layer_verify_compatible(why)) {
err = "verify: " + why + " (the verify window reproduces the default native decode path)";
return false;
}
if (!strata::kernels::fused_gr_supported(g.n_embd, g.hc, g.hc_lr) || ss.k != 10 || g.ssm_state_size != 128 ||
g.ssm_d_conv != 4) {
err = "verify: geometry differs from the artifact's";
return false;
}
if (le_ < 0) le_ = g.n_layers;
if (lb_ < 0 || lb_ >= le_ || le_ > g.n_layers || (lb_ > 0 && hand_in_ == nullptr) ||
(le_ < g.n_layers && hand_out_ == nullptr)) {
err = "verify: the stage's layer range or its hand-off buffers are wrong";
return false;
}
const WeightRef* wo = wt.find("output.weight");
if (wo == nullptr) { err = "verify: output.weight is missing"; return false; }
n_vocab_ = wo->ne1;
const strata::kernels::QsaShapes s = shapes_of(g);
cap_ = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s);
max_blocks_ = ss.qsa_states[ss.qsa_primary()].max_cells / s.idx_block + 2;
attn_scratch_floats_ = (int64_t) strata::kernels::qsa_decode_attn_scratch_floats(cap_, s);
const uint64_t T = (uint64_t) max_t, N = (uint64_t) g.n_embd, HC = (uint64_t) g.hc, K = (uint64_t) ss.k;
const uint64_t C = (uint64_t) g.ssm_conv_channels, ZV = (uint64_t) g.ssm_value_dim, HV = (uint64_t) g.ssm_v_heads;
const uint64_t NH = (uint64_t) g.n_head, HD = (uint64_t) g.head_dim, NKV = (uint64_t) g.n_head_kv;
const uint64_t IQ = (uint64_t) g.idx_q_heads, ID = (uint64_t) g.idx_key_dim;
const uint64_t nG = (uint64_t) g.n_gdn_layers(), nQ = (uint64_t) g.n_qsa_layers();
const uint64_t HS = (uint64_t) strata::kernels::NG_HIST * strata::kernels::NG_HC_DIM;
const uint64_t TS = (uint64_t) (s.idx_block - 1) * ID;
const int max_in = (int) std::max<uint64_t>(std::max<uint64_t>(N, ZV), NH * HD);
// ---- mapped staging
bool ok = mapped(T * 4, (void**) &h_tok_, (void**) &m_tok_) &&
mapped(T * strata::kernels::kStepCount * 4, (void**) &h_step_, (void**) &m_step_) &&
mapped(T * (NH + NKV + IQ) * 4, (void**) &h_pos_, (void**) &m_pos_) &&
mapped((2 + T) * 4 + 16, (void**) &h_commit_, (void**) &m_commit_) &&
mapped(T * N * 4, (void**) &h_ple_, (void**) &m_ple_) &&
mapped(T * 4 + 16, (void**) &h_out_, (void**) &m_out_) &&
mapped(T * N * 4, (void**) &h_x_, (void**) &m_x_) &&
mapped(T * K * 4, (void**) &h_ids_, (void**) &m_ids_) &&
mapped(T * K * 4, (void**) &h_w_, (void**) &m_w_) &&
mapped(64, (void**) &h_seq_, (void**) &m_seq_) &&
mapped(64, (void**) &h_flag_, (void**) &m_flag_) &&
mapped(64, (void**) &h_flagA_, (void**) &m_flagA_) &&
mapped(64, (void**) &h_flagB_, (void**) &m_flagB_) &&
mapped(T * K * N * 4, (void**) &h_ymiss_, (void**) &m_ymiss_);
if (!ok) { err = "verify: mapped staging allocation failed"; return false; }
// the GPU plan: counts(4) | start(cap+1) | dst(cap) | tok(cap) | pad | ptr(cap u64) | ptr2(cap u64) | start2(cap+1)
{
const int64_t cap = (int64_t) (T * K);
const int64_t i32 = 4 + (cap + 1) + cap + cap;
const int64_t ptr_off = (i32 + 1) & ~1ll;
plan_i32_ = ptr_off + 4 * cap + (cap + 1) + 1;
if (!mapped((size_t) plan_i32_ * 4 * 2 + 64, (void**) &h_plan_, (void**) &m_plan_)) {
err = "verify: mapped plan allocation failed";
return false;
}
sink_.counts = h_plan_;
sink_.start = h_plan_ + 4;
sink_.dst = sink_.start + cap + 1;
sink_.tok = sink_.dst + cap;
sink_.ptr = (unsigned long long*) (h_plan_ + ptr_off);
sink_.ptr2 = sink_.ptr + cap;
sink_.start2 = h_plan_ + ptr_off + 4 * cap;
sink_.cap = cap;
sink_.publish = &Verifier::publish_plan;
sink_.fetch = &Verifier::fetch_dma;
sink_.ctx = this;
}
// ---- the device arena: the same sequence counted, then carved
auto carve = [&](Bump& b) {
tok_ = b.take<int32_t>(T); step_ = b.take<int32_t>(T * strata::kernels::kStepCount);
pos_ = b.take<int32_t>(T * (NH + NKV + IQ)); commit_ = b.take<int32_t>(2 + T);
ple_ = b.take<float>(T * N); emb_ = b.take<float>(T * N); R_ = b.take<float>(T * HC * N);
mixed_ = b.take<float>(T * N); bo_ = b.take<float>(T * N);
inj_ = b.take<float>(T * HC); inj2_ = b.take<float>(T * HC);
lo_ = b.take<float>(T * (uint64_t) g.hc_lr); rs_ = b.take<float>(T * HC); xn_ = b.take<float>(T * HC * N);
xq_ = b.take<uint8_t>(strata::kernels::native_q8_1_bytes(max_in, (int) T));
qkv_L_ = b.take<float>(nG * T * C); h_L_ = b.take<float>(nG * T * C);
gate_L_ = b.take<float>(nG * T * HV); beta_L_ = b.take<float>(nG * T * HV);
z_ = b.take<float>(T * ZV); y_ = b.take<float>(T * ZV); y_dummy_ = b.take<float>(T * ZV);
qfull_ = b.take<float>(T * NH * 2 * HD); qcur_ = b.take<float>(T * NH * HD);
kcur_ = b.take<float>(T * NKV * HD); vcur_ = b.take<float>(T * NKV * HD);
idx_raw_L_ = b.take<float>(nQ * T * ID); qidx_ = b.take<float>(T * IQ * ID);
scores_ = b.take<float>(T * (uint64_t) max_blocks_); sel_ = b.take<int32_t>(T * (uint64_t) cap_);
attn_ = b.take<float>(T * NH * HD); attn32_ = b.take<float>(T * NH * HD);
attn_scratch_ = b.take<float>(T * (uint64_t) attn_scratch_floats_);
tail_snap_ = b.take<float>(nQ * TS);
logits_ = b.take<float>(T * (uint64_t) g.n_expert); w_ = b.take<float>(T * K); ids_ = b.take<int32_t>(T * K);
shared_ = b.take<float>(T * N); parts_ = b.take<float>(T * K * N); hit_out_ = b.take<float>(T * K * N);
hit_slot_ = b.take<int32_t>(T * K); hit_dst_ = b.take<int32_t>(T * K); hit_count_ = b.take<int32_t>(4);
plan_ = b.take<int32_t>(2 * ((uint64_t) plan_i32_ + 16));
staging_ = b.take<uint8_t>((uint64_t) kStagingBlobs * strata::kernels::cpu::expert_layout().max_blob);
hit_xq_ = b.take<uint8_t>(T * (N / 32) * 34); hit_xs_ = b.take<float>(T * (N / 32));
nat_xq_ = b.take<uint8_t>(T * (N / 32) * 36);
hit_scratch_ = b.take<uint8_t>(std::max<uint64_t>(
strata::kernels::moe_hit_grouped_scratch_bytes((int64_t) (T * K), g.n_embd, g.n_ff),
strata::kernels::native_expert_scratch_bytes((int64_t) (T * K), g.n_ff)));
head_mixed_ = b.take<float>(T * N); head_inj_ = b.take<float>(HC);
sh_bf16_ = b.take<uint16_t>(T * N); sh_gate_ = b.take<float>(T * (uint64_t) g.n_ff);
sh_up_ = b.take<float>(T * (uint64_t) g.n_ff); sh_g_ = b.take<float>(T + 4);
head_logits_ = b.take<float>(T * (uint64_t) n_vocab_);
hist_snap_ = b.take<float>(T * HS);
};
Bump count;
carve(count);
if (cudaMalloc(&arena_, count.used) != cudaSuccess) {
err = "verify: the device arena (" + std::to_string(count.used >> 20) + " MiB) does not fit";
return false;
}
cudaMemset(arena_, 0, count.used);
prof_on_ = std::getenv("STRATA_VERIFY_PROFILE") != nullptr;
if (prof_on_) {
const size_t np = (size_t) g.n_layers * kProfPer + 4;
if (cudaMalloc((void**) &prof_, np * 8) != cudaSuccess) { prof_on_ = false; prof_ = nullptr; cudaGetLastError(); }
else { cudaMemset(prof_, 0, np * 8); prof_h_.assign(np, 0); }
}
Bump real;
real.base = (uint8_t*) arena_;
carve(real);
sink_.staging = (unsigned long long) staging_;
sink_.staging_cap = kStagingBlobs;
(void) TS;
if (cudaStreamCreateWithFlags(&copy_, cudaStreamNonBlocking) != cudaSuccess) {
err = "verify: copy stream create failed";
return false;
}
if (cudaStreamCreateWithFlags(&cs_, cudaStreamNonBlocking) != cudaSuccess) {
err = "verify: stream create failed";
return false;
}
if (cudaEventCreateWithFlags(&commit_done_, cudaEventDisableTiming) != cudaSuccess) {
err = "verify: event create failed";
return false;
}
// E-6: a layer whose routed experts are all resident is planned on the device (STRATA_VERIFY_DEVICE_PLAN=1: on;
// exact, but neutral on RIBPC 1-2 GPUs: off by default)
{
const char* v = std::getenv("STRATA_VERIFY_DEVICE_PLAN");
device_plan_ = v != nullptr && std::atoi(v) != 0;
}
if (device_plan_) {
bool ok2 = cudaMalloc((void**) &skip_, 64) == cudaSuccess && cudaMemset(skip_, 0, 64) == cudaSuccess;
if (ok2 && hits.slot_off != nullptr && hits.n_slots > 0) {
ok2 = cudaMalloc((void**) &slot_off_d_, (size_t) hits.n_slots * sizeof(unsigned long long)) == cudaSuccess &&
cudaMemcpy(slot_off_d_, hits.slot_off, (size_t) hits.n_slots * sizeof(unsigned long long),
cudaMemcpyHostToDevice) == cudaSuccess;
}
if (!ok2) { cudaGetLastError(); device_plan_ = false; }
}
std::fprintf(stderr, "strata verify: window up to %d tokens, %.1f MiB of device buffers\n", max_t,
(double) count.used / 1048576.0);
return true;
}
const float* Verifier::final_R(int t) const { return R_ + (size_t) t * (size_t) (g_->hc * g_->n_embd); }
// ================================ THE WINDOW, AS CAPTURED ================================
//
// Plan v0.3 P6 (split window): with `groups_ == 2` the window's tokens are cut into two groups A = [0, T/2 up) and
// B = the rest, and the stream is ordered
//
// pre(0,A) pre(0,B) | post(0,A) pre(1,A) | post(0,B) pre(1,B) | post(1,A) pre(2,A) | ...
//
// so the CPU computes A's experts of layer l while the GPU runs B's mixer and router of layer l, and B's experts
// while the GPU combines A and runs A's layer l+1. B's mixer only needs A's mixer of the same layer (K/V, GDN
// state), never A's experts, so nothing waits that did not wait before. Every token's arithmetic is unchanged.
bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
using namespace strata::kernels;
const ModelGeometry& g = *g_;
const WeightTable& wt = *wt_;
SessionState& ss = *ss_;
const int64_t N = g.n_embd, HC = g.hc, K = ss.k, C = g.ssm_conv_channels, ZV = g.ssm_value_dim;
const int64_t HV = g.ssm_v_heads, HK = g.ssm_k_heads, NH = g.n_head, HD = g.head_dim, NKV = g.n_head_kv;
const int64_t IQ = g.idx_q_heads, ID = g.idx_key_dim, NE = g.n_expert, MT = max_t_;
const QsaShapes s = shapes_of(g);
const GrShapes gs{g.n_embd, g.hc, g.hc_lr};
const uint64_t gdn_floats = (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size +
(uint64_t) g.ssm_conv_channels * (g.ssm_d_conv - 1);
const int64_t HS = (int64_t) NG_HIST * NG_HC_DIM;
const int64_t TS = (s.idx_block - 1) * ID;
const bool ple_on = ss.ple.ready() && ple_stage();
auto Rt = [&](int t) { return R_ + (size_t) t * HC * N; };
const int G = (split_ && T >= 2) ? 2 : 1;
static const bool dec_batch = [] { const char* v = std::getenv("STRATA_DEC_BATCH"); return v == nullptr || std::atoi(v) != 0; }();
auto stamp = [&](int64_t l, int i, int grp) { if (prof_on_ && grp == 0) gpu_stamp(prof_, (int) (l * kProfPer + i), cs); };
const int tb_[2] = {0, (T + 1) / 2}, te_[2] = {G == 2 ? (T + 1) / 2 : T, T};
groups_[T] = G;
// ---- the window's inputs, from mapped staging
copy_i32_from_mapped(tok_, m_tok_, T, cs);
copy_i32_from_mapped(step_, m_step_, (int64_t) T * kStepCount, cs);
copy_i32_from_mapped(pos_, m_pos_, (int64_t) MT * (NH + NKV + IQ), cs);
// per-ROW positions of the K rows [t][NKV] and the indexer query rows [t][IQ] (for batched RoPE)
const int32_t* pos_k = pos_ + MT * NH;
const int32_t* pos_i = pos_ + MT * (NH + NKV);
if (ple_on) copy_from_mapped(ple_, m_ple_, (int64_t) T * N, cs);
// ---- the embeddings, broadcast to the hc streams - or, in a later stage of a layer split, the previous stage's
// residual, pending write and inject (see set_stage)
const int64_t HB = Verifier::handoff_floats(g);
if (lb_ > 0) {
for (int t = 0; t < T; ++t) {
copy_from_mapped(Rt(t), hand_in_ + (size_t) t * HB, HC * N, cs);
copy_from_mapped(bo_ + (size_t) t * N, hand_in_ + (size_t) t * HB + HC * N, N, cs);
copy_from_mapped(inj2_ + (size_t) t * HC, hand_in_ + (size_t) t * HB + HC * N + N, HC, cs);
}
} else if (const NativeEmbed* ne = native_embed()) { // plan v0.3 P6: the GGUF-form table
ne->gather_dev(tok_, T, emb_, cs);
broadcast_streams(emb_, R_, N, (int) HC, T, cs);
} else {
const WeightRef* w = wt.find("token_embd.weight");
if (w == nullptr || w->codebook_iq4nl || (w->code_bits != 2 && w->code_bits != 4 && w->code_bits != 8)) {
err = "verify: token_embd.weight is missing or not an S2/S4/S8 tensor";
return false;
}
const auto* codes = (const uint8_t*) w->data;
const auto* scales = (const float*) (codes + w->codes_bytes);
const auto* offsets = w->has_offset ? (const float*) (codes + w->codes_bytes + w->scales_bytes) : nullptr;
const uint64_t row_codes = (uint64_t) (w->ne0 / (8 / w->code_bits));
const uint64_t row_groups = (uint64_t) (w->ne0 / w->group_elems);
embedding_gather_dev(codes, scales, offsets, tok_, T, w->ne0, w->code_bits, w->code_bias, w->group_elems,
row_codes, row_groups, emb_, cs);
broadcast_streams(emb_, R_, N, (int) HC, T, cs);
}
// per-layer state indices (GDN and QSA layers are numbered separately)
std::vector<int64_t> gdn_idx((size_t) g.n_layers, -1), qsa_idx((size_t) g.n_layers, -1);
{
int64_t qi = 0, gi = 0;
for (int64_t l = 0; l < g.n_layers; ++l) {
if (is_qsa_layer(g, l)) qsa_idx[(size_t) l] = qi++;
else gdn_idx[(size_t) l] = gi++;
}
}
// ---------------------------------------------------------------- pre(l, group): up to the ring
auto pre = [&](int64_t l, int grp) -> bool {
const int tb = tb_[grp], te = te_[grp], n = te - tb;
stamp(l, 0, grp);
const LayerView v(wt, l);
const char* pfx[2] = {"hc_attn_", "hc_ffn_"};
const WeightRef *wn[2], *wd[2], *wu[2], *wi[2];
for (int h = 0; h < 2; ++h) {
wn[h] = need(v, (std::string(pfx[h]) + "norm.weight").c_str(), err);
wd[h] = need(v, (std::string(pfx[h]) + "down.weight").c_str(), err);
wu[h] = need(v, (std::string(pfx[h]) + "up.weight").c_str(), err);
wi[h] = need(v, (std::string(pfx[h]) + "inject.weight").c_str(), err);
if (!wn[h] || !wd[h] || !wu[h] || !wi[h]) return false;
}
// the previous layer's FFN write, folded into this layer's first read (a control vector after it has
// already applied it)
bool pending = l > 0 && !cvec().covers(l - 1);
if (l == 1 && ple_on) {
float* normalized = (float*) ((uint8_t*) ss.ple.scratch + ple_block_scratch_bytes());
for (int t = tb; t < te; ++t) {
gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs);
PleOut po;
po.normalized = normalized;
po.result = Rt(t);
try {
ple_block(ple_ + t * N, Rt(t), ss.ple.hist, ss.ple.w, po, ss.ple.scratch, cs);
ple_history_advance(ss.ple.hist, normalized, cs);
} catch (const std::exception& e) {
err = std::string("verify PLE: ") + e.what();
return false;
}
copy_from_mapped(hist_snap_ + (size_t) t * HS, ss.ple.hist, HS, cs);
}
pending = false;
}
auto gr_read_group = [&](int half, bool apply, float* inj_prev, float* inj_out) {
FusedGrArgs fa[kFusedGrMaxT];
for (int t = tb; t < te; ++t) {
FusedGrArgs& a = fa[t - tb];
a.R = Rt(t); a.R_out = Rt(t); a.apply = apply;
a.bo_prev = bo_ + t * N; a.inj_prev = inj_prev + t * HC;
a.w_norm = (const float*) wn[half]->data; a.w_down = (const uint16_t*) wd[half]->data;
a.w_up = (const uint16_t*) wu[half]->data; a.w_inject = (const uint16_t*) wi[half]->data;
a.eps = EPS; a.lo = lo_ + t * g.hc_lr; a.rs = rs_ + t * HC;
a.inject_out = inj_out + t * HC; a.mixed = mixed_ + t * N;
}
fused_gr_read_multi(fa, n, xn_ + (size_t) tb * HC * N, cs, (prof_on_ && grp == 0) ? prof_ : nullptr,
(int) (l * kProfPer + (half == 0 ? 27 : 30)));
};
gr_read_group(0, pending, inj2_, inj_);
stamp(l, 1, grp);
float* xm = mixed_ + tb * N;
try {
if (!is_qsa_layer(g, l)) {
// ======================= GDN =======================
const WeightRef *wqkv = need(v, "attn_qkv.weight", err), *wg = need(v, "attn_gate.weight", err),
*wout = need(v, "ssm_out.weight", err), *wa = need(v, "ssm_alpha.weight", err),
*wb = need(v, "ssm_beta.weight", err), *wc = need(v, "ssm_conv1d.weight", err),
*wnm = need(v, "ssm_norm.weight", err), *wdt = need(v, "ssm_dt.bias", err),
*wsa = need(v, "ssm_a", err);
if (!wqkv || !wg || !wout || !wa || !wb || !wc || !wnm || !wdt || !wsa) return false;
if (!native_of(wqkv, v.name("attn_qkv.weight"), err) || !native_of(wg, v.name("attn_gate.weight"), err) ||
!native_of(wout, v.name("ssm_out.weight"), err))
return false;
const int64_t gi = gdn_idx[(size_t) l];
float* state = ss.gdn_state + (size_t) (gi - ss.gdn_ord0) * gdn_floats;
float* conv = state + (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size;
float* qkv = qkv_L_ + (size_t) gi * MT * C;
float* hb = h_L_ + (size_t) gi * MT * C;
float* gate = gate_L_ + (size_t) gi * MT * HV;
float* beta = beta_L_ + (size_t) gi * MT * HV;
native_quantize_q8_1(xm, xq_, (int) N, n, cs);
native_mmvq(wqkv->native_type, wqkv->native_data, xq_, qkv + (size_t) tb * C, (int) N, (int) C, n, cs);
stamp(l, 2, grp);
gdn_conv_l2_multi(conv, qkv, (const float*) wc->data, hb, (int) C, (int) (2 * HK), EPS, n, cs, tb);
stamp(l, 3, grp);
gdn_ab_multi(xm, (const uint16_t*) wa->data, (const uint16_t*) wb->data, (const float*) wdt->data,
(const float*) wsa->data, gate + (size_t) tb * HV, beta + (size_t) tb * HV, (int) N, (int) HV,
n, cs);
stamp(l, 4, grp);
native_mmvq(wg->native_type, wg->native_data, xq_, z_ + (size_t) tb * ZV, (int) N, (int) ZV, n, cs);
stamp(l, 5, grp);
// the recurrence from the untouched state over tokens [0, te); outputs only for this group's
gdn_step_norm_multi(state, hb, (int) C, gate, beta, z_, (const float*) wnm->data, EPS, y_, (int) HK,
(int) HV, te, nullptr, cs, tb);
stamp(l, 6, grp);
native_quantize_q8_1(y_ + (size_t) tb * ZV, xq_, (int) ZV, n, cs);
native_mmvq(wout->native_type, wout->native_data, xq_, bo_ + tb * N, (int) ZV, (int) N, n, cs);
} else {
// ======================= QSA =======================
const int64_t qi = qsa_idx[(size_t) l];
const QsaState& st = ss.qsa_states[qi];
const WeightRef *wik = need(v, "indexer.k_proj.weight", err), *wq = need(v, "attn_q.weight", err),
*wk = need(v, "attn_k.weight", err), *wv = need(v, "attn_v.weight", err),
*wo = need(v, "attn_output.weight", err), *wiq = need(v, "indexer.q_proj.weight", err),
*wqn = need(v, "attn_q_norm.weight", err), *wkn = need(v, "attn_k_norm.weight", err),
*wiqn = need(v, "indexer.q_norm.weight", err), *wikn = need(v, "indexer.k_norm.weight", err);
if (!wik || !wq || !wk || !wv || !wo || !wiq || !wqn || !wkn || !wiqn || !wikn) return false;
if (!native_of(wq, v.name("attn_q.weight"), err) || !native_of(wk, v.name("attn_k.weight"), err) ||
!native_of(wv, v.name("attn_v.weight"), err) || !native_of(wo, v.name("attn_output.weight"), err))
return false;
auto norm_rope = [&](float* data, const WeightRef* norm, int rows, int cols, const int32_t* pos) {
if (native_qsa_enabled()) native_qsa_rms_norm_weighted(data, (const float*) norm->data, data, cols, rows, EPS, cs);
else rms_norm_weighted(data, (const float*) norm->data, rows, cols, EPS, cs);
if (native_rope_enabled()) native_rope_apply(data, data, rows, cols, (int) s.n_rot, rope_scaling(), pos, cs);
else rope_neox_apply(data, data, rows, cols, (int) s.n_rot, st.cos_tab, st.sin_tab, pos, cs);
};
float* idx_raw = idx_raw_L_ + (size_t) qi * MT * ID;
// the per-token GEMVs / norms / RoPEs / copies of this layer as one launch over the
// window's rows each - row-wise identical arithmetic (STRATA_DEC_BATCH=0: token by token)
const bool qb = dec_batch && n > 1 && native_qsa_enabled() && native_rope_enabled() && !st.kv_q4;
native_quantize_q8_1(xm, xq_, (int) N, n, cs);
if (qb) bf16_gemv_fp32_mmvf_multi(mixed_ + tb * N, N, (const uint16_t*) wik->data, idx_raw + tb * ID, ID, N, ID, n, cs);
else for (int t = tb; t < te; ++t)
bf16_gemv_fp32_mmvf(mixed_ + t * N, (const uint16_t*) wik->data, idx_raw + t * ID, (int) N, (int) ID, cs);
stamp(l, 7, grp);
native_mmvq(wk->native_type, wk->native_data, xq_, kcur_ + tb * NKV * HD, (int) N, (int) (NKV * HD), n, cs);
native_mmvq(wv->native_type, wv->native_data, xq_, vcur_ + tb * NKV * HD, (int) N, (int) (NKV * HD), n, cs);
if (qb) norm_rope(kcur_ + tb * NKV * HD, wkn, (int) (n * NKV), (int) HD, pos_k + tb * NKV);
else for (int t = tb; t < te; ++t) norm_rope(kcur_ + t * NKV * HD, wkn, (int) NKV, (int) HD, pos_ + t * NH);
if (st.kv_rot) { // K and V rotated before they are stored (kv_q4.hpp)
fwht256_inplace_cuda(kcur_ + tb * NKV * HD, (int64_t) n * NKV, cs);
fwht256_inplace_cuda(vcur_ + tb * NKV * HD, (int64_t) n * NKV, cs);
} else if (st.kv_hybrid) { // K8V4: only V is rotated
fwht256_inplace_cuda(vcur_ + tb * NKV * HD, (int64_t) n * NKV, cs);
}
stamp(l, 8, grp);
if (grp == 0) copy_from_mapped(tail_snap_ + (size_t) qi * TS, st.idx_tail, TS, cs);
for (int t = tb; t < te; ++t) {
const int32_t* step_t = step_ + t * kStepCount;
if (st.kv_hybrid) { // K8V4: the unused half's lanes folded onto the used pool (layer.cpp)
kv_append_q8_step(st.k_q, st.k_q, st.k_scale, st.k_scale, st.page_table, step_t,
kcur_ + t * NKV * HD, kcur_ + t * NKV * HD, s, cs, nullptr);
kv_append_q4_step(st.v_q4, st.v_q4, st.page_table, step_t, vcur_ + t * NKV * HD,
vcur_ + t * NKV * HD, s, cs, nullptr);
} else if (st.kv_q4)
kv_append_q4_step(st.k_q4, st.v_q4, st.page_table, step_t, kcur_ + t * NKV * HD,
vcur_ + t * NKV * HD, s, cs, &st.host);
else 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, &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, &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)
native_qsa_indexer_append(idx_raw + t * ID, step_ + t * kStepCount + kStepPos, 0,
(const float*) wikn->data, EPS, ib, s, st.max_cells,
rope_scaling(), cs);
stamp(l, 9, grp);
native_mmvq(wq->native_type, wq->native_data, xq_, qfull_ + tb * NH * 2 * HD, (int) N, (int) (NH * 2 * HD),
n, cs);
if (qb) {
if (cudaMemcpy2DAsync(qcur_ + tb * NH * HD, (size_t) HD * 4, qfull_ + tb * NH * 2 * HD, (size_t) HD * 2 * 4,
(size_t) HD * 4, (size_t) (n * NH), cudaMemcpyDeviceToDevice, cs) != cudaSuccess) {
err = "verify: the q/gate split failed";
return false;
}
norm_rope(qcur_ + tb * NH * HD, wqn, (int) (n * NH), (int) HD, pos_ + tb * NH);
if (st.kv_rot) fwht256_inplace_cuda(qcur_ + tb * NH * HD, (int64_t) n * NH, cs); // <Hq, Hk> = <q, k>
bf16_gemv_fp32_mmvf_multi(mixed_ + tb * N, N, (const uint16_t*) wiq->data, qidx_ + tb * IQ * ID, IQ * ID,
N, IQ * ID, n, cs);
norm_rope(qidx_ + tb * IQ * ID, wiqn, (int) (n * IQ), (int) ID, pos_i + tb * IQ);
} else {
for (int t = tb; t < te; ++t) {
float* qc = qcur_ + t * NH * HD;
if (cudaMemcpy2DAsync(qc, (size_t) HD * 4, qfull_ + t * NH * 2 * HD, (size_t) HD * 2 * 4,
(size_t) HD * 4, (size_t) NH, cudaMemcpyDeviceToDevice, cs) != cudaSuccess) {
err = "verify: the q/gate split failed";
return false;
}
norm_rope(qc, wqn, (int) NH, (int) HD, pos_ + t * NH);
if (st.kv_rot) fwht256_inplace_cuda(qc, NH, cs); // <Hq, Hk> = <q, k>
}
for (int t = tb; t < te; ++t) {
float* qx = qidx_ + t * IQ * ID;
bf16_gemv_fp32_mmvf(mixed_ + t * N, (const uint16_t*) wiq->data, qx, (int) N, (int) (IQ * ID), cs);
norm_rope(qx, wiqn, (int) IQ, (int) ID, pos_ + t * NH);
}
}
stamp(l, 10, grp);
qsa_block_scores(st.idx_pooled, st.idx_dead, qidx_ + tb * IQ * ID, step_ + tb * kStepCount, n, max_blocks_,
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);
stamp(l, 11, grp);
// 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);
stamp(l, 12, grp);
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);
stamp(l, 13, grp);
if (st.kv_rot || st.kv_hybrid) fwht256_inplace_cuda(attn_ + tb * NH * HD, (int64_t) n * NH, cs); // back: H^-1 = H
if (qb) native_qsa_gate_apply(attn_ + tb * NH * HD, qfull_ + tb * NH * 2 * HD, attn32_ + tb * NH * HD,
(int) (n * NH), (int) HD, cs);
else
for (int t = tb; t < te; ++t) {
if (native_qsa_enabled())
native_qsa_gate_apply(attn_ + t * NH * HD, qfull_ + t * NH * 2 * HD, attn32_ + t * NH * HD,
(int) NH, (int) HD, cs);
else
qsa_gate_apply_f32(attn_ + t * NH * HD, qfull_ + t * NH * 2 * HD, s, attn32_ + t * NH * HD, cs);
}
stamp(l, 14, grp);
native_quantize_q8_1(attn32_ + tb * NH * HD, xq_, (int) (NH * HD), n, cs);
native_mmvq(wo->native_type, wo->native_data, xq_, bo_ + tb * N, (int) (NH * HD), (int) N, n, cs);
}
} catch (const std::exception& e) {
err = "verify layer " + std::to_string(l) + ": " + e.what();
return false;
}
stamp(l, 16, grp);
gr_read_group(1, true, inj_, inj2_);
// the window's rows routed in 2 launches (one router GEMV reading the weight once, one
// top-10) instead of 2 per token; every row's arithmetic is the single-token call's (STRATA_DEC_BATCH=0: old)
const WeightRef* w_router = v.get("ffn_gate_inp.weight");
if (dec_batch && n > 1 && w_router != nullptr && native_router_enabled() && NE == 512 && K == 10) {
try {
bf16_gemv_fp32_mmvf_multi(mixed_ + tb * N, N, (const uint16_t*) w_router->data, logits_ + tb * NE, NE, N,
NE, n, cs);
native_router_top10_multi(logits_ + tb * NE, ids_ + tb * K, w_ + tb * K, n, cs);
} catch (const std::exception& e) { err = "verify router: " + std::string(e.what()); return false; }
} else
for (int t = tb; t < te; ++t) {
MoEBuffers mb = ss.moe;
mb.logits = logits_ + t * NE; mb.ids = ids_ + t * K; mb.weights = w_ + t * K;
if (!moe_route(wt, g, l, K, mb, mixed_ + t * N, cs, err, nullptr)) return false;
}
if (device_plan_) // E-6: every routed expert resident: this group's plan without the host
resident_plan(ids_ + tb * K, n * (int) K, (int) K, hits_.d_res + l * g.n_expert, (int) g.n_expert,
hits_.cache_base, slot_off_d_, (long long) hits_.blob,
plan_ + (size_t) grp * (size_t) (plan_i32_ + 16), (long long) max_t_ * K, skip_ + grp,
(uint32_t) ((l - lb_) * G + grp + 1), cs);
doorbell_publish(xm, ids_ + tb * K, w_ + tb * K, (int64_t) n * N, (int64_t) n * K, m_x_ + tb * N,
m_ids_ + tb * K, m_w_ + tb * K, m_seq_, cs);
stamp(l, 17, grp);
{
const WeightRef *wgi = need(v, "ffn_gate_inp_shexp.weight", err), *wsg = need(v, "ffn_gate_shexp.weight", err),
*wsu = need(v, "ffn_up_shexp.weight", err), *wsd = need(v, "ffn_down_shexp.weight", err);
if (!wgi || !wsg || !wsu || !wsd) return false;
if (!native_of(wsg, v.name("ffn_gate_shexp.weight"), err) || !native_of(wsu, v.name("ffn_up_shexp.weight"), err) ||
!native_of(wsd, v.name("ffn_down_shexp.weight"), err))
return false;
NativeSharedWeights nsw;
nsw.gate_type = wsg->native_type; nsw.gate_data = wsg->native_data;
nsw.up_type = wsu->native_type; nsw.up_data = wsu->native_data;
nsw.down_type = wsd->native_type; nsw.down_data = wsd->native_data;
nsw.q8_1 = xq_;
if (dec_batch) f32_to_bf16_bulk(mixed_ + tb * N, sh_bf16_ + tb * N, (int64_t) n * N, cs); // contiguous rows
else for (int t = tb; t < te; ++t) f32_to_bf16_bulk(mixed_ + t * N, sh_bf16_ + t * N, N, cs);
try {
shared_expert_multi(n, xm, sh_bf16_ + tb * N, nsw, (const uint16_t*) wgi->data, sh_gate_ + (size_t) tb * g.n_ff,
sh_up_ + (size_t) tb * g.n_ff, sh_g_ + tb, shared_ + tb * N, N, g.n_ff, cs);
} catch (const std::exception& e) {
err = std::string("verify shared expert: ") + e.what();
return false;
}
}
if (strata::kernels::cpu::expert_layout().native)
quantize_q8_1_rows(xm, n, N, nat_xq_ + (size_t) tb * (N / 32) * 36, cs);
else
quantize_q8_0_scaled(xm, hit_xq_ + (size_t) tb * (N / 32) * 34, hit_xs_ + (size_t) tb * (N / 32), (int64_t) n * N, cs);
stamp(l, 18, grp);
return true;
};
// ---------------------------------------------------------------- post(l, group): experts, combine
auto post = [&](int64_t l, int grp) -> bool {
const int tb = tb_[grp], te = te_[grp], n = te - tb;
const uint32_t ring = (uint32_t) ((l - lb_) * G + grp + 1);
const int64_t cap = (int64_t) n * K, capx = (int64_t) max_t_ * K;
int32_t* pl = plan_ + (size_t) grp * (size_t) (plan_i32_ + 16);
if (device_plan_) { // E-6: skipped when the device planned this group (all its experts resident)
wait_flag_ge_or(m_flagA_, ring, skip_ + grp, cs);
copy_i32_from_mapped_unless(pl, m_plan_ + (size_t) grp * (size_t) plan_i32_, plan_i32_, skip_ + grp, ring, cs);
} else {
wait_flag_ge(m_flagA_, ring, cs); // the pool published this group's GPU plan
copy_i32_from_mapped(pl, m_plan_ + (size_t) grp * (size_t) plan_i32_, plan_i32_, cs);
}
stamp(l, 19, grp);
const int32_t* p_counts = pl;
const int32_t* p_start = pl + 4;
const int32_t* p_dst = p_start + capx + 1;
const int32_t* p_tok = p_dst + capx;
const int64_t ptr_off = ((4 + (capx + 1) + 2 * capx) + 1) & ~1ll;
const unsigned long long* p_ptr = (const unsigned long long*) (pl + ptr_off);
const unsigned long long* p_ptr2 = p_ptr + capx;
const int32_t* p_start2 = pl + ptr_off + 4 * capx;
float* hit_out = hit_out_ + (size_t) tb * K * N;
const auto& lay = strata::kernels::cpu::expert_layout();
// plan v0.3 P6: the VRAM groups now; the PCIe groups once the copy engine has landed them in staging.
// `gy`: the native launch's groups side by side (0: cap, one block row per possible group).
auto grouped = [&](const unsigned long long* gp, const int32_t* gs, const int32_t* gn, int64_t gy) {
if (lay.native) {
// the layer's GGUF formats (i-quant gate/up, Q2_0 / IQ4_NL down)
const auto& f = lay.fmt[(size_t) l];
const NativeExpertLayout L = native_expert_layout(f.gu_type, f.d_type, f.n_embd, f.n_ff);
native_expert_grouped(L, gp, gs, gn, p_dst, p_tok, cap, cap,
nat_xq_ + (size_t) tb * (N / 32) * 36, hit_scratch_, hit_out, cs, gy);
} else {
moe_grouped_s2(gp, gs, gn, p_dst, p_tok, cap, cap, hit_xq_ + (size_t) tb * (N / 32) * 34,
hit_xs_ + (size_t) tb * (N / 32), hit_scratch_, hit_out, cs);
}
};
grouped(p_ptr, p_start, p_counts, 0);
stamp(l, 20, grp);
if (device_plan_) wait_flag_ge_or(m_flagB_, ring, skip_ + grp, cs);
else wait_flag_ge(m_flagB_, ring, cs); // the PCIe share is in staging (DMA) or mapped
if (sink_.pcie_mode == 2) { // stage it with a copy kernel, then point at staging
const int64_t per = G == 2 ? kStagingBlobs / 2 : kStagingBlobs;
uint8_t* stage = staging_ + (size_t) (grp * per) * lay.max_blob;
fetch_blobs(p_ptr2, p_counts + 2, stage, (int64_t) lay.blob_bytes(l), (int) per, cs);
rebase_ptrs((unsigned long long*) p_ptr2, p_counts + 2, stage, (int64_t) lay.blob_bytes(l), cs);
}
stamp(l, 21, grp);
// the PCIe share is pcie_frac of the misses: a few groups when the cache is cold, usually none (always none at
// pcie_frac 0), so its launch is kPcieGroupRows block rows striding over the groups, not cap of them
grouped(p_ptr2, p_start2, p_counts + 2, kPcieGroupRows);
stamp(l, 22, grp);
if (device_plan_) { // no CPU share when the device planned the group: its rows are zeros
wait_flag_ge_or(m_flag_, ring, skip_ + grp, cs);
copy_or_zero_from_mapped(parts_ + (size_t) tb * K * N, m_ymiss_ + (size_t) tb * K * N, (long long) n * K * N,
skip_ + grp, ring, cs);
} else {
wait_flag_ge(m_flag_, ring, cs); // the CPU's share is in the mapped rows
stamp(l, 23, grp);
if (dec_batch) // only the CPU rows cross PCIe (p_dst[0, counts[1]) = the GPU's own rows)
copy_rows_from_mapped(parts_ + (size_t) tb * K * N, m_ymiss_ + (size_t) tb * K * N, (int64_t) n * K, N,
p_dst, p_counts + 1, cs);
else
copy_from_mapped(parts_ + (size_t) tb * K * N, m_ymiss_ + (size_t) tb * K * N, (int64_t) n * K * N, cs);
}
moe_hit_add(parts_ + (size_t) tb * K * N, hit_out, p_dst, p_counts + 1, cap, N, cs);
if (dec_batch && n > 1 && native_moe_combine_enabled()) { // one launch for the window's rows
try {
native_moe_combine_multi(parts_ + (size_t) tb * K * N, w_ + tb * K, shared_ + tb * N, bo_ + tb * N, N, K, n, cs);
} catch (const std::exception& e) { err = "verify combine: " + std::string(e.what()); return false; }
} else
for (int t = tb; t < te; ++t) {
MoEBuffers mb = ss.moe;
mb.weights = w_ + t * K; mb.shared = shared_ + t * N;
if (!moe_combine_parts(g, l, K, mb, parts_ + (size_t) t * K * N, bo_ + t * N, cs, err)) return false;
}
stamp(l, 24, grp);
if (l == g.n_layers - 1) {
for (int t = tb; t < te; ++t) gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs);
if (cvec().covers(l)) cvec_apply(Rt(tb), l, n, HC * N, nullptr, 0, nullptr, 0, false, cs);
} else if (cvec().covers(l)) {
cvec_apply(Rt(tb), l, n, HC * N, bo_ + tb * N, N, inj2_ + tb * HC, HC, true, cs);
}
return true;
};
for (int grp = 0; grp < G; ++grp)
if (!pre(lb_, grp)) return false;
for (int64_t l = lb_; l < le_; ++l)
for (int grp = 0; grp < G; ++grp) {
if (!post(l, grp)) return false;
if (l + 1 < le_ && !pre(l + 1, grp)) return false;
}
if (le_ < g.n_layers) { // a layer split's earlier stage: hand the residual on, no head
for (int t = 0; t < T; ++t) {
copy_from_mapped(hand_out_ + (size_t) t * HB, Rt(t), HC * N, cs);
copy_from_mapped(hand_out_ + (size_t) t * HB + HC * N, bo_ + (size_t) t * N, N, cs);
copy_from_mapped(hand_out_ + (size_t) t * HB + HC * N + N, inj2_ + (size_t) t * HC, HC, cs);
}
return true;
}
// ---- the head, T columns, and the argmax of each
stamp(g.n_layers, 0, 0);
{
const WeightRef *hn = wt.find("output_hc_norm.weight"), *hd = wt.find("output_hc_down.weight"),
*hu = wt.find("output_hc_up.weight");
if (!hn || !hd || !hu) { err = "verify: an output_hc_* weight is missing"; return false; }
for (int t = 0; t < T; ++t) {
BlockBuffers bb = ss.block;
bb.R = Rt(t);
bb.mixed = head_mixed_ + t * N;
if (head_ != nullptr && head_->loaded()) {
if (!lm_head_mix(wt, g, bb, cs, err)) return false;
} else if (!lm_head(wt, g, bb, head_logits_ + (size_t) t * n_vocab_, cs, err)) {
return false;
}
}
if (head_ != nullptr && head_->loaded()) {
try {
native_quantize_q8_1(head_mixed_, xq_, (int) N, T, cs);
native_mmvq(head_->type(), head_->weights(), xq_, head_logits_, (int) N, (int) n_vocab_, T, cs);
} catch (const std::exception& e) {
err = std::string("verify head: ") + e.what();
return false;
}
}
// Greedy, the default, is recorded here as before (no extra launch or sync per window). A request that
// samples or penalizes is sampled again host-side after the replay (run()) with its own parameters and a
// fresh draw counter: a captured sampler would bake them in and replay the same draws forever.
SamplerParams sp;
sp.greedy = true;
sp.temperature = 0.0f;
sample_tokens(head_logits_, T, (int) n_vocab_, nullptr, 0, sp, m_out_, cs);
}
stamp(g.n_layers, 1, 0);
return true;
}
std::string Verifier::profile_report() {
if (!prof_on_ || prof_windows_ == 0) return std::string();
static const char* names[kProfPer] = {"-", "hc-read0", "q8+qkv/q-idx gemv", "conv", "ab", "z", "rec", "q8+kv-idx",
"k/v+norm-rope", "kv+idx append", "q+q-idx", "scores+topk", "kv-resolve",
"attention", "gate", "", "out-proj", "hc-read1+router", "shared+quant",
"waitA", "VRAM hits", "waitB", "PCIe grp", "waitCPU", "copy+combine",
"(gap)", "head", " hc0 norm", " hc0 down", " hc0 up", "", "", ""};
std::string out;
char b[80];
double total = 0;
for (int k = 0; k < 2; ++k) {
out += k == 0 ? " GDN layers:" : " | QSA layers:";
for (int i = 0; i < kProfPer; ++i) {
if (prof_sum_[k][i] <= 0) continue;
total += prof_sum_[k][i];
std::snprintf(b, sizeof b, " %s %.2f", names[i], prof_sum_[k][i] / 1e6 / (double) prof_windows_);
out += b;
}
}
std::snprintf(b, sizeof b, " | total %.2f ms/window over %lld windows", total / 1e6 / (double) prof_windows_, (long long) prof_windows_);
out += b;
for (auto& r : prof_sum_) for (double& d : r) d = 0;
prof_windows_ = 0;
return out;
}
bool Verifier::capture(int T, std::string& err) {
if (exec_[T] != nullptr) return true;
if (cudaStreamBeginCapture(cs_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) {
err = "verify: begin capture failed";
return false;
}
std::string rerr;
const bool ok = record_window(T, cs_, rerr);
cudaGraph_t graph = nullptr;
const cudaError_t ce = cudaStreamEndCapture(cs_, &graph);
if (!ok) {
if (graph) cudaGraphDestroy(graph);
err = rerr;
return false;
}
if (ce != cudaSuccess) {
err = std::string("verify: end capture: ") + cudaGetErrorString(ce);
return false;
}
#if !defined(STRATA_USE_HIP) // a CUDA debug listing (node types, kernel names)
if (std::getenv("STRATA_VERIFY_NODES") != nullptr) { // what the window graph holds
size_t nn = 0;
cudaGraphGetNodes(graph, nullptr, &nn);
std::vector<cudaGraphNode_t> nodes(nn);
cudaGraphGetNodes(graph, nodes.data(), &nn);
std::map<std::string, int> kinds;
for (cudaGraphNode_t nd : nodes) {
cudaGraphNodeType ty;
cudaGraphNodeGetType(nd, &ty);
std::string name = "type" + std::to_string((int) ty);
if (ty == cudaGraphNodeTypeKernel) {
cudaKernelNodeParams kp{};
if (cudaGraphKernelNodeGetParams(nd, &kp) == cudaSuccess) {
#if CUDART_VERSION >= 12030 // cudaFuncGetName arrived in CUDA 12.3
const char* fn = nullptr;
if (cudaFuncGetName(&fn, kp.func) == cudaSuccess && fn) name = fn;
#endif
}
} else if (ty == cudaGraphNodeTypeMemcpy) name = "memcpy";
else if (ty == cudaGraphNodeTypeMemset) name = "memset";
++kinds[name];
}
std::vector<std::pair<int, std::string>> v;
for (auto& [k2, c] : kinds) v.push_back({c, k2});
std::sort(v.rbegin(), v.rend());
std::fprintf(stderr, "strata verify: the %d-token window graph has %zu nodes:", T, nn);
for (size_t i = 0; i < v.size() && i < 40; ++i) std::fprintf(stderr, " %d x %.60s;", v[i].first, v[i].second.c_str());
std::fprintf(stderr, "\n");
}
#endif
const cudaError_t ie = cudaGraphInstantiate(&exec_[T], graph, 0);
cudaGraphDestroy(graph);
if (ie != cudaSuccess) {
err = std::string("verify: instantiate: ") + cudaGetErrorString(ie);
return false;
}
const cudaError_t ue = cudaGraphUpload(exec_[T], cs_);
const cudaError_t us = cudaStreamSynchronize(cs_);
std::fprintf(stderr, "strata verify: captured the %d-token window (upload %s, sync %s)\n", T,
cudaGetErrorString(ue), cudaGetErrorString(us));
return true;
}
bool Verifier::capture_commit(std::string& err) {
if (commit_exec_ != nullptr) return true;
using namespace strata::kernels;
const ModelGeometry& g = *g_;
SessionState& ss = *ss_;
const QsaShapes s = shapes_of(g);
const int64_t C = g.ssm_conv_channels, HV = g.ssm_v_heads, ID = g.idx_key_dim, MT = max_t_;
const uint64_t gdn_floats = (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size +
(uint64_t) g.ssm_conv_channels * (g.ssm_d_conv - 1);
const int64_t TS = (s.idx_block - 1) * ID;
const int64_t HS = (int64_t) NG_HIST * NG_HC_DIM;
if (cudaStreamBeginCapture(cs_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) {
err = "verify: begin commit capture failed";
return false;
}
bool ok = true;
try {
copy_i32_from_mapped(commit_, m_commit_, 2 + MT, cs_);
int64_t qsa_index = 0, gdn_index = 0;
for (int64_t l = 0; l < lb_; ++l) (is_qsa_layer(g, l) ? qsa_index : gdn_index) += 1;
for (int64_t l = lb_; l < le_ && ok; ++l) {
const LayerView v(*wt_, l);
if (!is_qsa_layer(g, l)) {
const WeightRef* wnm = need(v, "ssm_norm.weight", err);
if (!wnm) { ok = false; break; }
float* state = ss.gdn_state + (size_t) (gdn_index - ss.gdn_ord0) * gdn_floats;
float* conv = state + (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size;
const float* qkv = qkv_L_ + (size_t) gdn_index * MT * C;
gdn_conv_commit(conv, qkv, (int) C, commit_, cs_);
gdn_step_norm_multi(state, h_L_ + (size_t) gdn_index * MT * C, (int) C, gate_L_ + (size_t) gdn_index * MT * HV,
beta_L_ + (size_t) gdn_index * MT * HV, z_, (const float*) wnm->data, EPS, y_dummy_,
(int) g.ssm_k_heads, (int) HV, (int) MT, commit_, cs_);
++gdn_index;
} else {
const QsaState& st = ss.qsa_states[qsa_index];
const WeightRef* wikn = need(v, "indexer.k_norm.weight", err);
if (!wikn) { ok = false; break; }
copy_from_mapped(st.idx_tail, tail_snap_ + (size_t) qsa_index * TS, TS, cs_);
const QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos};
for (int64_t t = 0; t < MT; ++t)
native_qsa_indexer_append(idx_raw_L_ + (size_t) (qsa_index * MT + t) * ID, commit_ + 2 + t, 0,
(const float*) wikn->data, EPS, ib, s, st.max_cells,
rope_scaling(), cs_);
++qsa_index;
}
}
if (ok && ss.ple.ready() && ple_stage()) copy_indexed(ss.ple.hist, hist_snap_, HS, commit_ + 1, HS, cs_);
} catch (const std::exception& e) {
err = std::string("verify commit: ") + e.what();
ok = false;
}
cudaGraph_t graph = nullptr;
const cudaError_t ce = cudaStreamEndCapture(cs_, &graph);
if (!ok) {
if (graph) cudaGraphDestroy(graph);
return false;
}
if (ce != cudaSuccess || cudaGraphInstantiate(&commit_exec_, graph, 0) != cudaSuccess) {
if (graph) cudaGraphDestroy(graph);
err = std::string("verify: commit capture: ") + cudaGetErrorString(ce);
return false;
}
cudaGraphDestroy(graph);
return true;
}
bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, void* user, int32_t* out,
std::string& err) {
using namespace strata::kernels;
const OnDevice on_device(device_);
if (T < 1 || T > max_t_) { err = "verify: window size out of range"; return false; }
if (released_.load()) { err = "verify: an earlier window never finished on the GPU (#267); restart the engine"; return false; }
const ModelGeometry& g = *g_;
SessionState& ss = *ss_;
if (pos0 + T > ss.qsa_states[ss.qsa_primary()].max_cells) { err = "verify: the window runs past the context"; return false; }
if (!capture(T, err) || !capture_commit(err)) return false;
VDBG("captured; staging\n");
const Clock::time_point t0 = Clock::now();
const QsaShapes s = shapes_of(g);
for (int t = 0; t < T; ++t) {
h_tok_[t] = tokens[t];
qsa_step_fill(h_step_ + t * kStepCount, pos0 + t, s);
for (int64_t h = 0; h < g.n_head; ++h) h_pos_[t * g.n_head + h] = (int32_t) (pos0 + t);
int32_t* pk = h_pos_ + (size_t) max_t_ * g.n_head;
int32_t* pi = pk + (size_t) max_t_ * g.n_head_kv;
for (int64_t h = 0; h < g.n_head_kv; ++h) pk[t * g.n_head_kv + h] = (int32_t) (pos0 + t);
for (int64_t h = 0; h < g.idx_q_heads; ++h) pi[t * g.idx_q_heads + h] = (int32_t) (pos0 + t);
}
if (ss.ple.ready() && ple_stage()) {
uint32_t rows[kVerifyMaxT * PLE_N_HEADS];
int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]};
for (int t = 0; t < T; ++t) {
ngram_rows(&tokens[t], prev, 1, ss.ple.consts, rows + t * PLE_N_HEADS);
prev[0] = prev[1];
prev[1] = tokens[t];
}
if (!ss.ple.table->gather_batch(rows, (size_t) T, h_ple_, err)) return false;
}
*(volatile uint32_t*) h_seq_ = 0;
*(volatile uint32_t*) h_flag_ = 0;
*(volatile uint32_t*) h_flagA_ = 0;
*(volatile uint32_t*) h_flagB_ = 0;
std::atomic_thread_fence(std::memory_order_seq_cst);
last_t_ = T;
last_pos0_ = pos0;
for (int t = 0; t < T; ++t) last_tokens_[t] = tokens[t];
ms_host += ms_since(t0);
VDBG("staged; launching\n");
const cudaError_t le = cudaGraphLaunch(exec_[T], cs_);
if (le != cudaSuccess) { err = std::string("verify: launch: ") + cudaGetErrorString(le); return false; }
(void) cudaStreamQuery(cs_);
VDBG("launched\n");
volatile uint32_t* const seq = h_seq_;
volatile uint32_t* const flag = h_flag_;
const int G = groups_[T] > 0 ? groups_[T] : 1;
const int gtb[2] = {0, (T + 1) / 2}, gte[2] = {G == 2 ? (T + 1) / 2 : T, T};
const int64_t steps = (le_ - lb_) * G;
const bool test_stall = g_test_stall > 0 && windows + 1 == g_test_stall; // #267 test hook (off: false)
for (int64_t k = 0; k < steps; ++k) {
const int64_t l = lb_ + k / G;
const int grp = (int) (k % G);
const uint32_t want = (uint32_t) (k + 1);
const Clock::time_point a = Clock::now();
auto last_flush = a;
uint32_t spins = 0;
progress_at("verify window: waiting for the GPU to reach layer", l);
while (*seq < want) {
_mm_pause();
if ((++spins & 1023u) != 0) continue;
const auto now = Clock::now();
if (now - last_flush > std::chrono::microseconds(2000)) {
last_flush = now;
const cudaError_t q = cudaStreamQuery(cs_);
if (q != cudaErrorNotReady && *seq < want) {
err = "verify: layer " + std::to_string(l) + " never rang (" +
(q == cudaSuccess ? std::string("graph finished") : std::string(cudaGetErrorString(q))) + ")";
return false;
}
}
if (now - a > std::chrono::seconds(20)) {
// #267: the caller ends the engine; no spin kernel may outlive it
err = "verify: timed out at layer " + std::to_string(l) + released_note(release_gpu_waits(5000));
return false;
}
}
const Clock::time_point b = Clock::now();
VDBG("layer %lld rang\n", (long long) l);
cur_layer_ = want - 1;
set_plan_slot(grp);
const int tb = gtb[grp], n = gte[grp] - gtb[grp];
progress_at("verify window: the CPU experts of layer", l);
if (pool != nullptr)
pool(user, h_x_ + (size_t) tb * g.n_embd, h_ids_ + (size_t) tb * ss.k, n, ss.k,
h_ymiss_ + (size_t) tb * ss.k * g.n_embd, l);
VDBG("layer %lld served\n", (long long) l);
progress_tick();
std::atomic_thread_fence(std::memory_order_seq_cst);
_mm_sfence();
if (*(volatile uint32_t*) h_flagA_ != want) { // the pool did not publish a plan: an empty one
sink_.counts[0] = 0;
sink_.counts[1] = 0;
sink_.counts[2] = 0;
sink_.start[0] = 0;
sink_.start2[0] = 0;
std::atomic_thread_fence(std::memory_order_seq_cst);
*(volatile uint32_t*) h_flagA_ = want;
raise_flag(h_flagB_, want);
}
if (!(test_stall && k + 1 == steps)) *flag = want;
ms_wait += std::chrono::duration<double, std::milli>(b - a).count();
ms_pool += ms_since(b);
}
progress_at("verify window: waiting for the GPU to finish the window (flags A/B/M raised)", (int64_t) T);
// #267: a window the GPU never finishes (a spin kernel that never sees its flag) holds the host here; the stall
// watchdog then releases every verifier's GPU waits (release_live_verifiers) before it ends the engine, so no
// spin kernel outlives the process - the case that left Windows GPUs "lost" until a power cycle. The wait
// itself stays a blocking sync: a cudaStreamQuery poll here cost IQ3_S ~3% decode (a core calling the driver
// beside the expert workers).
const cudaError_t se = cudaStreamSynchronize(cs_);
if (se != cudaSuccess) { err = std::string("verify: ") + cudaGetErrorString(se); return false; }
progress_at("verify window: waiting for the expert copies", (int64_t) T);
cudaStreamSynchronize(copy_); // no host function of this window may raise flag B in the next one
if (prof_on_ && G == 1) { // the window's GPU stage stamps
cudaMemcpy(prof_h_.data(), prof_, prof_h_.size() * 8, cudaMemcpyDeviceToHost);
const int64_t L = g.n_layers;
auto at = [&](int64_t l, int i) { return prof_h_[(size_t) (l * kProfPer + i)]; };
// D8: the derived columns below read stamps the hc-read kernels write themselves or the next
// layer's first. A slot no kernel stamped is 0 and its unsigned difference wrapped to ~1e19 ns -
// which is why (gap)/head/hc0 printed ~1e14 ms per window. A missing or out-of-order stamp now
// contributes nothing.
const auto gap = [](unsigned long long to, unsigned long long from) {
return (from != 0 && to != 0 && to >= from) ? (double) (to - from) : 0.0;
};
for (int64_t l = 0; l < L; ++l) {
const int kind = is_qsa_layer(g, l) ? 1 : 0;
unsigned long long prev = at(l, 0);
for (int i = 1; i <= 24; ++i) {
const unsigned long long x = at(l, i);
if (x == 0 || x < prev) continue;
prof_sum_[kind][i] += (double) (x - prev);
prev = x;
}
if (l + 1 < L) prof_sum_[kind][25] += gap(at(l + 1, 0), at(l, 24));
const double dn = gap(at(l, 27), at(l, 0)), dd = gap(at(l, 28), at(l, 27)), du = gap(at(l, 1), at(l, 28));
if (dn > 0 && dd > 0 && du > 0) { // the split exists: show it split, not twice
prof_sum_[kind][27] += dn; // hc-read0: norm
prof_sum_[kind][28] += dd; // down
prof_sum_[kind][29] += du; // up (through both halves, as before)
prof_sum_[kind][1] -= gap(at(l, 1), at(l, 0)); // (hc-read0 shown split)
}
}
prof_sum_[0][26] += gap(at(L, 1), at(L, 0));
++prof_windows_;
}
// ---- a sampled or penalized request: the head's sampling again, host-side so its parameters are this call's
// own (a captured kernel would replay the same draws forever). Row t's draw is Philox(seed, pos0 + t): tied to
// the POSITION it samples, not to how the text was cut into windows, so a seed replays the same text whatever
// the drafts were. Exact: a rejected row's draw is discarded, and no kept decision depends on a reused draw.
if (le_ < g.n_layers) { // a layer split's earlier stage: the hand-off is written (synced above)
++windows;
return next_ == nullptr || next_->run(T, tokens, pos0, pool, next_user_, out, err);
}
const bool sampled = !sampling_.greedy && sampling_.temperature > 0.0f;
if (head_sampling_ && (sampled || hist_d_ != nullptr)) {
SamplerParams sp = sampling_;
sp.counter = (uint64_t) pos0;
sample_tokens(head_logits_, T, (int) n_vocab_, hist_d_, hist_len_, sp, m_out_, cs_);
if (cudaStreamSynchronize(cs_) != cudaSuccess) { // m_out_ is the mapped h_out_: synced, it is readable
err = "verify: the head sampling failed";
return false;
}
}
for (int t = 0; t < T; ++t) out[t] = ((volatile int32_t*) h_out_)[t];
if (static const bool dbg = std::getenv("STRATA_DBG_NAN") != nullptr; dbg) { // debug: the first non-finite head
static bool reported = false;
if (!reported) {
std::vector<float> h((size_t) T * (size_t) n_vocab_);
cudaMemcpy(h.data(), head_logits_, h.size() * 4, cudaMemcpyDeviceToHost);
for (int t = 0; t < T && !reported; ++t) {
int64_t bad = 0;
for (int64_t v = 0; v < n_vocab_; ++v) bad += !std::isfinite(h[(size_t) t * n_vocab_ + v]);
if (bad) {
reported = true;
std::fprintf(stderr, "strata dbg: verify window at position %lld, row %d: %lld of %lld logits non-finite "
"(token out %d)\n", (long long) pos0, t, (long long) bad, (long long) n_vocab_, out[t]);
}
}
}
}
VDBG("window done\n");
++windows;
progress_at("decode");
progress_beat();
return true;
}
void Verifier::set_plan_slot(int grp) {
const int64_t cap = sink_.cap;
int32_t* base = h_plan_ + (size_t) grp * (size_t) plan_i32_;
const int64_t i32 = 4 + (cap + 1) + cap + cap;
const int64_t ptr_off = (i32 + 1) & ~1ll;
sink_.counts = base;
sink_.start = base + 4;
sink_.dst = sink_.start + cap + 1;
sink_.tok = sink_.dst + cap;
sink_.ptr = (unsigned long long*) (base + ptr_off);
sink_.ptr2 = sink_.ptr + cap;
sink_.start2 = base + ptr_off + 4 * cap;
const int G = groups_[last_t_] > 0 ? groups_[last_t_] : 1;
const int64_t per = G == 2 ? kStagingBlobs / 2 : kStagingBlobs;
sink_.staging = (unsigned long long) (staging_ + (size_t) (grp * per) * strata::kernels::cpu::expert_layout().max_blob);
sink_.staging_cap = per;
}
// Flag B only rises: a host function of an earlier layer may run after a later layer already raised it directly.
void Verifier::raise_flag(uint32_t* flag, uint32_t value) {
volatile long* f = (volatile long*) flag;
#if defined(_WIN32)
long cur = *f;
while ((uint32_t) cur < value) {
const long prev = _InterlockedCompareExchange(f, (long) value, cur);
if (prev == cur) break;
cur = prev;
}
#else
uint32_t cur = __atomic_load_n((uint32_t*) flag, __ATOMIC_SEQ_CST);
while (cur < value && !__atomic_compare_exchange_n((uint32_t*) flag, &cur, value, false, __ATOMIC_SEQ_CST, __ATOMIC_SEQ_CST)) {}
#endif
}
// Plan v0.3 P6: the PCIe share by DMA. The copy engine moves the blobs while the CPU computes its own share and the
// GPU its VRAM experts; a host function raises flag B when they have landed (the graph waits for it before the PCIe
// groups). Staging is split between the two token groups of a split window.
void Verifier::fetch_dma(void* ctx, const uint8_t* const* src, int n, size_t bytes) {
Verifier* v = (Verifier*) ctx;
const uint32_t want = v->cur_layer_ + 1;
if (n <= 0) { raise_flag(v->h_flagB_, want); return; }
uint8_t* stage = (uint8_t*) v->sink_.staging; // this group's half in a split window
for (int i = 0; i < n; ++i) cudaMemcpyAsync(stage + (size_t) i * bytes, src[i], bytes, cudaMemcpyHostToDevice, v->copy_);
FlagSet& fs = v->flag_sets_[v->cur_layer_ % (sizeof v->flag_sets_ / sizeof v->flag_sets_[0])];
fs.flag = v->h_flagB_;
fs.value = want;
cudaLaunchHostFunc(v->copy_, [](void* p) { FlagSet* s = (FlagSet*) p; raise_flag(s->flag, s->value); }, &fs);
}
void Verifier::publish_plan(void* ctx) {
Verifier* v = (Verifier*) ctx;
_mm_sfence();
*(volatile uint32_t*) v->h_flagA_ = v->cur_layer_ + 1;
}
bool Verifier::window_logprobs(const int32_t* targets, int T, int64_t pos0, int32_t extra_id, std::FILE* out,
std::string& err) {
if (next_ != nullptr) return next_->window_logprobs(targets, T, pos0, extra_id, out, err);
const OnDevice on_device(device_);
if (head_logits_ == nullptr || n_vocab_ <= 0 || T <= 0 || targets == nullptr || out == nullptr) {
err = "window_logprobs: no head logits for this window";
return false;
}
// run() synchronized cs_ before returning, so the head of this window is complete
std::vector<float> h((size_t) T * (size_t) n_vocab_);
if (cudaMemcpy(h.data(), head_logits_, h.size() * sizeof(float), cudaMemcpyDeviceToHost) != cudaSuccess) {
err = "window_logprobs: the head logits copy failed";
return false;
}
for (int t = 0; t < T; ++t) {
const float* row = h.data() + (size_t) t * (size_t) n_vocab_;
const int32_t tgt = targets[t];
if (tgt < 0 || (int64_t) tgt >= n_vocab_) continue;
int64_t top = 0;
for (int64_t v = 1; v < n_vocab_; ++v)
if (row[v] > row[top]) top = v;
const double maxv = row[top];
const bool has_extra = extra_id >= 0 && (int64_t) extra_id < n_vocab_;
double sum = 0.0, sum_without = 0.0; // the second skips extra_id: no cancellation when it holds ~all mass
for (int64_t v = 0; v < n_vocab_; ++v) {
const double e = std::exp((double) row[v] - maxv);
sum += e;
if (v != (int64_t) extra_id) sum_without += e;
}
const double lse = maxv + std::log(sum);
const double extra = has_extra ? (double) row[extra_id] - lse : NAN;
const double without = has_extra && tgt != extra_id && sum_without > 0.0
? (double) row[tgt] - (maxv + std::log(sum_without)) : NAN;
std::fprintf(out, "%lld\t%d\t%.9f\t%lld\t%.9f\t%d\t%.9f\t%.9f", (long long) (pos0 + t), (int) tgt,
(double) row[tgt] - lse, (long long) top, maxv - lse, (int) (top == (int64_t) tgt), extra,
without);
// STRATA_LOGPOS_TOPK=K: the K most likely tokens and their log-probabilities too (`id:logprob`), for a
// top-k comparison with another engine on the same tokens (docs/UNSLOTH_Q4.md)
static const int topk = [] {
const char* v = std::getenv("STRATA_LOGPOS_TOPK");
return v != nullptr ? std::max(0, std::min(256, std::atoi(v))) : 0;
}();
if (topk > 0) {
std::vector<int32_t> order((size_t) n_vocab_);
for (int64_t v = 0; v < n_vocab_; ++v) order[(size_t) v] = (int32_t) v;
std::partial_sort(order.begin(), order.begin() + topk, order.end(),
[&](int32_t a, int32_t b) { return row[a] > row[b]; });
for (int j = 0; j < topk; ++j)
std::fprintf(out, "\t%d:%.6f", order[(size_t) j], (double) row[order[(size_t) j]] - lse);
}
std::fprintf(out, "\n");
}
std::fflush(out);
return true;
}
namespace { bool g_commit_async = false; }
void Verifier::set_commit_async(bool on) { g_commit_async = on && std::getenv("STRATA_COMMIT_SYNC") == nullptr; }
bool Verifier::commit(int n_keep, std::string& err) {
const OnDevice on_device(device_);
if (n_keep < 1 || n_keep > last_t_) { err = "verify: commit count out of range"; return false; }
const Clock::time_point t0 = Clock::now();
h_commit_[0] = n_keep;
h_commit_[1] = n_keep - 1;
for (int t = 0; t < max_t_; ++t) h_commit_[2 + t] = t < n_keep ? (int32_t) (last_pos0_ + t) : -1;
std::atomic_thread_fence(std::memory_order_seq_cst);
const cudaError_t le = cudaGraphLaunch(commit_exec_, cs_);
if (le != cudaSuccess) { err = std::string("verify: commit launch: ") + cudaGetErrorString(le); return false; }
// set_commit_async: no wait here - the next window runs on the same stream after it, and the drafter (its own
// stream) reads only this window's final rows and its own K/V. h_commit_ is next written after the next window's
// results are read, i.e. after this graph has run. Everything else waits on commit_done_ (wait_commit).
if (!g_commit_async || next_ != nullptr) {
const cudaError_t se = cudaStreamSynchronize(cs_);
if (se != cudaSuccess) { err = std::string("verify: commit: ") + cudaGetErrorString(se); return false; }
} else {
const cudaError_t re = cudaEventRecord(commit_done_, cs_);
if (re != cudaSuccess) { err = std::string("verify: commit event: ") + cudaGetErrorString(re); return false; }
commit_pending_ = true;
}
if (ple_stage()) // stages that share one session must advance it once
for (int t = 0; t < n_keep; ++t) {
ss_->ple_prev[0] = ss_->ple_prev[1];
ss_->ple_prev[1] = last_tokens_[t];
}
ms_commit += ms_since(t0);
return next_ == nullptr || next_->commit(n_keep, err);
}
bool Verifier::wait_commit(std::string& err) {
if (commit_pending_) {
const OnDevice on_device(device_);
commit_pending_ = false;
const cudaError_t se = cudaEventSynchronize(commit_done_);
if (se != cudaSuccess) { err = std::string("verify: commit: ") + cudaGetErrorString(se); return false; }
}
return next_ == nullptr || next_->wait_commit(err);
}
bool Verifier::copy_logits(int t, float* host) const {
if (next_ != nullptr) return next_->copy_logits(t, host); // a layer split: the head is on the last stage
if (head_logits_ == nullptr || host == nullptr || t < 0 || n_vocab_ <= 0) return false;
return cudaMemcpy(host, head_logits_ + (size_t) t * (size_t) n_vocab_, (size_t) n_vocab_ * sizeof(float),
cudaMemcpyDeviceToHost) == cudaSuccess;
}
} // namespace strata::core