Niko1221/Strata
1
0 Fork 0
mirror of https://github.com/Niko1221/Strata.git synced 2026-11-07 18:35:04 -10:00
Code Issues Packages Projects Releases 1 Wiki Activity

perf(core): wire multi-token kernels into Verifier and MtpDrafter CUDA graphs

Browse Source 
- verify.cpp / verify.hpp:
  - Dispatch gr_read_multi, gr_write_multi, native_qsa_rms_norm_rope,
    kv_append_q8_steps, kv_append_q4_steps, and
    native_qsa_indexer_append_steps across verify/batch tokens in
    Verifier::record_window and Verifier::capture_commit.
  - Add runtime env kill switches STRATA_NO_MULTI_GR=1,
    STRATA_NO_NORM_ROPE=1, and STRATA_NO_BATCH_KV_STEP=1 for isolated A/B
    bisection without recompiling.
  - Preserve 0.1.39's hit_out + moe_hit_add on the !all_resident_ path and
    synchronous --layer-split commit/stage ordering by default (deee447).
- mtp.cpp:
  - Dispatch gr_read_multi, gr_write_multi, native_qsa_rms_norm_rope,
    kv_append_q8_steps, and kv_append_q4_steps inside
    MtpDrafter::record_forward (honoring STRATA_NO_MULTI_GR,
    STRATA_NO_NORM_ROPE, and STRATA_NO_BATCH_KV_STEP).
This commit is contained in:
Stuart Chapin authored and root committed 2026-10-04 10:57:04 -07:00
1 parent b7222cc076
commit 4327807fc3
3 files changed
+3661 -3515
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

include/strata/core/verify.hpp
+365 -364
View File
@@ -1,364 +1,365 @@
// include/strata/core/verify.hpp - plan v0.3 P6: the speculative VERIFY window.
//
// T tokens at consecutive positions p0 .. p0+T-1 - the last accepted token and T-1 drafts - go through all 48
// layers in ONE captured graph, and the head's argmax is produced for every one of them. Token t's argmax is
// what plain greedy decode would produce after token t, BIT FOR BIT: every kernel here is either the single-token
// kernel applied per token, or a multi-token kernel whose per-token arithmetic is the single-token kernel's
// (multi-column MMVQ in exact mode, the T-token GDN kernels, the per-token hit activation, the multi-token CPU
// expert rows). So a draft is accepted exactly when greedy decode would have produced it.
//
// What the window costs is the dense weights read ONCE for T tokens and the union of the T tokens' missed
// experts on the CPU (measured on decode traces: 1.75x one token's misses for T=2, 2.4x for 3, 3.05x for 4).
//
// STATE. The window appends K/V and indexer keys for all T positions and leaves the GDN state untouched.
// `commit(n_keep)` then makes the first `n_keep` tokens permanent: the GDN conv history and recurrent state are
// advanced by replaying those tokens from inputs the window stored, the indexer's key tail is restored from a
// snapshot and the accepted keys re-appended (a rejected key can land in a slot the current block still needs),
// and the PLE history is set to its snapshot after token n_keep-1. K/V cells past the accepted prefix are simply
// overwritten when those positions are processed again, before any query can read them.
//
// Requires the default native decode configuration (native projections, fused GR, fused GDN, fast attention and
// selection, native indexer) and a profile-filled VRAM expert tier with its residency table on the device.
#pragma once
#include <cstdio>
#include "strata/core/expert_source.hpp"
#include "strata/core/layer.hpp"
#include "strata/core/session.hpp"
#include "strata/kernels/sampler.hpp"
#include <cuda_runtime.h>
#include <atomic>
#include <chrono>
#include <map>
#include <cstdint>
#include <string>
#include <vector>
namespace strata::core {
class NativeHead;
class RemoteExpertOpt;
/// The CPU pool for a window: x_f (n_tok, n_embd), ids (n_tok, k) -> out (n_tok * k, n_embd), hit rows zeroed.
using PoolMultiFn = void (*)(void* user, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k, float* out,
int64_t layer);
struct VerifyHits {
const int32_t* d_res = nullptr; ///< device [n_layers * n_expert] slot or -1
const int32_t* h_res = nullptr; ///< host [n_layers * n_expert] slot or -1 (when all >= 0 in stage: zero-doorbell)
const uint8_t* cache_base = nullptr; ///< slot 0 of the VRAM expert arena
const uint64_t* slot_off = nullptr; ///< E-6: host per-slot offsets when slots differ in size (null: slot * blob)
int64_t n_slots = 0; ///< E-6: how many (for the device copy)
int64_t blob = 0;
};
class Verifier {
public:
Verifier() = default;
~Verifier();
Verifier(const Verifier&) = delete;
Verifier& operator=(const Verifier&) = delete;
void set_remote_expert_opt(RemoteExpertOpt* opt) { remote_opt_ = opt; } // before init/capture
/// The watchdog's view of the window in flight (issue #31): the layer, the GPU's sequence, the flags.
void diag(std::FILE* f) const;
/// #649 (STRATA_VERIFY_TRACE=1, opt-in): the last windows' host-side handshake events and the GPU's breadcrumbs
/// (in mapped memory, so they read while the GPU hangs). Nothing when the trace is off.
void trace_dump(std::FILE* f) const;
/// #267: raise every flag the window's spin kernels wait on past any ring (UINT32_MAX), so a window the GPU
/// cannot finish drains instead of staying resident, then wait up to `timeout_ms` for its streams. For the
/// paths that give up on the engine (a timed-out window, the serve watchdog): the window then ran on whatever
/// the flags guarded, so this verifier refuses every later window. True when the streams finished.
bool release_gpu_waits(int timeout_ms);
/// `max_t` <= kVerifyMaxT. `head` may be null (the canonical head is then run per token).
bool init(const WeightTable& wt, const ModelGeometry& g, SessionState& ss, const VerifyHits& hits,
const NativeHead* head, int max_t, std::string& err);
/// One window: `tokens[0..T)` at positions pos0.., the pool served per layer; `out[t]` = argmax after token t.
/// The PLE rows are gathered here from `ss.ple_prev` and the tokens. Captures the T-token graph on first use.
bool run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, void* user, int32_t* out, std::string& err);
/// Diagnostics: row `t` of the last window's head logits (n_vocab floats) to the host. Valid after run().
bool copy_logits(int t, float* host) const;
int64_t vocab() const { return next_ ? next_->vocab() : n_vocab_; }
/// The sampling the verify window's head applies (temperature / top_p / top_k / seed). Set per
/// request; greedy by default. The sampling itself runs OUTSIDE the captured graph - its
/// parameters would otherwise be baked forever - so this can change between requests freely.
void set_sampling(const strata::kernels::SamplerParams& sp) {
sampling_ = sp; // row t of a window at pos0 draws Philox(seed, pos0 + t): see run()
if (next_) next_->set_sampling(sp);
}
/// The penalty histories for `sampling_.penalty_last_n`: ONE ROW PER WINDOW ROW, T rows of `history_len`
/// int32 slots at that stride (`strata::kernels::penalty_rows` builds them), most recent token LAST, unused
/// front slots -1 (the kernel reads only the tail window). Row t follows the window's drafts 1..t - staging
/// row 0 alone (before 0.1.19) left the drafted rows with unwritten histories. Null disables the penalties
/// entirely - the neutral run's sampling call is byte-for-byte what it was. The engine re-uploads the rows
/// before every window; the buffer must hold kVerifyMaxT rows and stay alive across the request.
void set_history(const int32_t* history, int history_len) {
hist_d_ = history;
hist_len_ = history_len;
if (next_) next_->set_history(history, history_len);
}
/// Off: `run` skips the request's head sampling and `out` is the recorded greedy pick. For windows whose
/// picks are discarded - a prompt read through windows commits every token - so they cost no sampler launch
/// or sync and never read a history staged for another position.
void set_head_sampling(bool on) { head_sampling_ = on; if (next_) next_->set_head_sampling(on); }
/// LAYER SPLIT (multi-GPU): this verifier runs layers [layer_begin, layer_end) of every window. A stage that
/// does not start at layer 0 takes its residual from `handoff_in` instead of embedding the tokens; a stage that
/// does not end at the last layer writes its residual to `handoff_out` and has no head. The hand-off holds,
/// per token, the residual R (hc x n_embd), the last layer's pending write bo (n_embd) and inject (hc): the next
/// stage folds that write into its first read exactly as the unsplit window does, so the split is bit-exact.
/// Both pointers must be device-visible (mapped pinned memory, portable when the stages are on two devices).
/// Set before `init`. Default: the whole model, no hand-off.
void set_stage(int64_t layer_begin, int64_t layer_end, const float* handoff_in, float* handoff_out) {
lb_ = layer_begin; le_ = layer_end; hand_in_ = handoff_in; hand_out_ = handoff_out;
}
/// The next stage: `run` and `commit` continue into it (its pool calls get `next_user`); sampling settings
/// and `final_R` are the last stage's.
void set_next(Verifier* next, void* next_user) { next_ = next; next_user_ = next_user; }
/// floats per token in a hand-off buffer
static int64_t handoff_floats(const ModelGeometry& g) { return (int64_t) g.hc * g.n_embd + g.n_embd + g.hc; }
/// Keep the first `n_keep` (1..T) tokens of the last window; advances `ss.ple_prev` by them.
bool commit(int n_keep, std::string& err);
// ================================ SEVERAL SEQUENCES IN ONE WINDOW ================================
//
// A batch window holds S INDEPENDENT sequences, one token each: row s is slot s, at slot s's own position,
// reading and writing slot s's own state (GDN recurrence and conv history, QSA K/V and indexer, PLE history),
// which lives in `slots[s]` - a session carved like this verifier's own (same layer range, same max_cells).
// Everything that is per row already (hyper-connections, dense projections, router, shared expert, the
// routed experts on the GPU and the CPU, the head) runs once over the S rows, so the weights are read once
// per window for all the sequences. Row s's arithmetic is the single-token window's, so a slot's greedy
// tokens are its solo greedy tokens (modulo the multi-token CPU kernel choice: STRATA_IQ_MT_MIN=1).
//
// Greedy only, no drafts (MTP) in a batch window. `init_slots` once after `init` (S <= max_t); a layer
// split's stages each get their own sessions, and run_slots/commit_slots continue into the next stage.
bool init_slots(const std::vector<SessionState*>& slots, std::string& err);
int n_slots() const { return (int) slots_.size(); }
/// One batch window over slots [0, S): tokens[s] at positions pos[s]; out[s] = the greedy pick after it.
bool run_slots(int S, const int32_t* tokens, const int64_t* pos, PoolMultiFn pool, void* user, int32_t* out,
std::string& err);
/// The same over the S slots `rows` (row t is slot rows[t], any distinct slots in any order): the slots not
/// listed are not touched, so an idle slot keeps its state (a finished conversation it may continue later).
bool run_slot_rows(const int* rows, int S, const int32_t* tokens, const int64_t* pos, PoolMultiFn pool, void* user,
int32_t* out, std::string& err);
/// Keep every row of the last batch window: each slot's state advances by its one token.
bool commit_slots(std::string& err);
// ---- The stages of a layer split as a PIPELINE. A batch window over the slot GROUP
// [base, base + S) is launched on ONE stage with its commit right behind it on the stage's stream (a batch window
// keeps every row, so the commit needs no host decision), and the host serves the rings of every stage that has
// a window in flight from one thread (batch_poll does not block). Stage k can then run group g while stage k+1
// runs group g-1. Rows of group `base` use hand-off rows [base, base + S), so groups never share a hand-off row.
bool batch_launch(int base, int S, const int32_t* tokens, const int64_t* pos, std::string& err);
/// 1 = this stage's window and commit are done (the last stage's picks are in batch_out), 0 = still running,
/// -1 = an error (err). Serves every layer that has rung so far.
int batch_poll(PoolMultiFn pool, void* user, std::string& err);
bool batch_busy() const { return b_running_; }
/// A slot's sampling (temperature / top_p / top_k / min_p / seed; penalties are not applied in batch windows):
/// its row is drawn again on the last stage with Philox(seed, position), as a solo window draws it. Greedy by
/// default. Set on the first stage, it reaches the last.
void set_slot_sampling(int slot, const strata::kernels::SamplerParams& sp) {
if (slot >= 0 && slot < (int) slot_sp_.size()) slot_sp_[(size_t) slot] = sp;
if (next_) next_->set_slot_sampling(slot, sp);
}
const int32_t* batch_out() const { return b_out_; }
bool last_stage() const { return g_ != nullptr && le_ == g_->n_layers; }
/// commit() returns without waiting for its graph (a single-GPU session sets it): the next window follows it on
/// the same stream and the drafter reads nothing it writes, so it overlaps the draft. Whoever reads or writes
/// the session from another stream or the host afterwards (a new request, a checkpoint, a snapshot, the prompt
/// path, the end of a run) calls wait_commit() first. STRATA_COMMIT_SYNC=1 keeps the wait.
static void set_commit_async(bool on);
/// Waits for the last commit graph when commit() did not (an event recorded after it, not the whole device);
/// false with `err` when it failed. Free when nothing is pending.
bool wait_commit(std::string& err);
/// Measurement hook (STRATA_LOGPOS): after run(), write one line per row t of the last window's head -
/// "pos target logprob top top_logprob hit extra_logprob target_logprob_without_extra" - where row t is the
/// distribution at pos0 + t, targets[t] is the token at pos0 + t + 1, extra_logprob is the log-probability of
/// `extra_id` in the same row, and the last column is the target's log-probability in that row renormalized
/// over every token but `extra_id` (both nan when they do not apply). A layer split's earlier stage forwards
/// to the stage that holds the head.
bool window_logprobs(const int32_t* targets, int T, int64_t pos0, int32_t extra_id, std::FILE* out,
std::string& err);
/// Token t's residual after the last layer, (hc, n_embd) on the device, valid until the next `run`.
const float* final_R(int t) const;
const float* final_R_all() const { return next_ ? next_->final_R_all() : R_; }
/// The GPU plan the pool writes each layer (VRAM hits + the PCIe share of the misses); give it to the
/// dispatch (`ExpertDispatch::plan`) before the first `run`.
GpuPlanSink* plan_sink() { return &sink_; }
/// Plan v0.3 P6: split the window into two token groups and pipeline the CPU experts of one with the GPU work
/// of the other (default on). Set before the first `run`.
void set_split(bool on) { split_ = on; }
/// Plan v0.3 P6: how the PCIe share of the misses reaches the GPU: 0 = DMA into staging (the copy engine works
/// beside the CPU; best when the CPU is compute-bound, the i-quants), 1 = the grouped kernel reads the mapped
/// arena directly, 2 = a copy kernel stages it inside the graph (no API calls on the pool's thread; best when
/// the CPU is RAM-bound, Q2_0). Set before the first `run`.
void set_pcie_mode(int mode) { sink_.pcie_mode = mode; }
/// the pool never plans a PCIe share (--pcie-frac 0): the window skips that path. Before the first run.
double ms_wait = 0, ms_pool = 0, ms_host = 0, ms_commit = 0;
int64_t windows = 0;
/// STRATA_VERIFY_PROFILE=1 - GPU stage times of the windows since the last call (ms per
/// window), as one line; empty when off.
std::string profile_report();
private:
RemoteExpertOpt* remote_opt_ = nullptr;
bool capture(int T, std::string& err);
// batch windows (see init_slots)
std::vector<SessionState*> slots_;
bool batch_rec_ = false; ///< record_window is capturing a batch window
int row_base_ = 0; ///< ... its hand-off rows start here (a pipeline group's own rows)
int brow_[8] = {}; ///< ... and row t is slot brow_[t]
bool last_batch_ = false; ///< the last run was a batch window (set_plan_slot: one group)
std::map<uint64_t, cudaGraphExec_t> exec_bm_, commit_bm_; ///< key: batch_key(rows, S, hand-off base)
int last_rows_[8] = {}; ///< the slots of the last batch window's rows
static uint64_t batch_key(const int* rows, int S, int hbase) {
uint64_t k = (uint64_t) hbase << 40 | (uint64_t) S << 32;
for (int t = 0; t < S; ++t) k |= (uint64_t) (rows[t] & 15) << (4 * t);
return k;
}
// batch_launch / batch_poll
bool b_running_ = false;
int64_t b_k_ = 0, b_steps_ = 0;
std::chrono::steady_clock::time_point b_last_;
int32_t b_out_[8] = {};
std::vector<strata::kernels::SamplerParams> slot_sp_;
bool sample_rows(int S, std::string& err); ///< the sampled slots' rows of the last batch window
int32_t* h_commitb_ = nullptr; int32_t* m_commitb_ = nullptr; // per slot [1, 0, pos, -1 ..], stride 2 + max_t
int32_t* commitb_ = nullptr;
float* tail_snap_b_ = nullptr; ///< per (slot, QSA layer) indexer tail snapshot
void* arena_b_ = nullptr;
int64_t last_pos_b_[8] = {};
bool capture_batch(const int* rows, int S, int hbase, std::string& err);
void collect_profile(); ///< STRATA_VERIFY_PROFILE: add the last window's stamps to prof_sum_
bool capture_commit_batch(const int* rows, int S, int hbase, std::string& err);
bool stage_batch(const int* rows, int S, int hbase, const int32_t* tokens, const int64_t* pos, std::string& err);
strata::kernels::SamplerParams sampling_ = [] {
strata::kernels::SamplerParams s;
s.greedy = true;
s.temperature = 0.0f;
return s;
}(); ///< greedy by default; per-request via set_sampling
const int32_t* hist_d_ = nullptr; ///< penalty-history row (set_history); null = no penalties apply
int hist_len_ = 0;
bool head_sampling_ = true; ///< set_head_sampling
int device_ = -1; ///< the device `init` ran on: run/commit switch to it (layer split)
std::atomic<bool> released_{false}; ///< #267: release_gpu_waits ran (maybe on the watchdog thread): no more windows
bool all_resident_ = false; ///< 100% of experts in [lb_, le_) resident in VRAM: zero-doorbell graph
bool device_plan_ = false; ///< E-6: resident-only layers planned on the device (STRATA_VERIFY_DEVICE_PLAN)
uint32_t* skip_ = nullptr; ///< E-6: per group, the ring whose plan the device built (0: the host's)
unsigned long long* slot_off_d_ = nullptr; ///< E-6: the slot offsets on the device
int64_t lb_ = 0, le_ = -1; ///< set_stage: the layers this verifier runs (-1: to the last)
const float* hand_in_ = nullptr;
float* hand_out_ = nullptr;
Verifier* next_ = nullptr;
void* next_user_ = nullptr;
bool ple_stage() const { return lb_ <= 1 && 1 < le_; } ///< holds layer 1, where the PLE block runs
void stage_inputs(int T, const int32_t* tokens, int64_t pos0);
bool staged_ = false;
bool copy_used_ = false;
bool capture_commit(std::string& err);
bool record_window(int T, cudaStream_t cs, std::string& err);
// #649: STRATA_VERIFY_TRACE=1 - a host event ring (trace_ev) and GPU breadcrumbs: the profiler's stamp points,
// per layer and token group, written to mapped memory (null when the trace is off)
void trace_ev(const char* what, int64_t step, int64_t layer, int64_t aux) const;
unsigned long long* trace_h_ = nullptr;
unsigned long long* trace_m_ = nullptr;
size_t trace_n_ = 0;
static constexpr int kProfPer = 33; // stamps per layer (32 left the hc-read second
// half's up-stamp at slot 32 = the next layer's slot 0: D8)
bool prof_on_ = false;
unsigned long long* prof_ = nullptr; // device: n_layers * kProfPer + 4 stamps
std::vector<unsigned long long> prof_h_;
double prof_sum_[2][kProfPer] = {}; // [GDN / QSA layers][stage]
int64_t prof_windows_ = 0;
const WeightTable* wt_ = nullptr;
const ModelGeometry* g_ = nullptr;
SessionState* ss_ = nullptr;
VerifyHits hits_;
const NativeHead* head_ = nullptr;
int max_t_ = 0;
int last_t_ = 0;
int64_t last_pos0_ = 0;
int32_t last_tokens_[8] = {};
int64_t n_vocab_ = 0;
cudaStream_t cs_ = nullptr;
cudaStream_t sh_cs_ = nullptr;
cudaEvent_t ev_fork_ = nullptr, ev_join_ = nullptr;
cudaGraphExec_t exec_[9] = {};
cudaGraphExec_t commit_exec_ = nullptr;
// mapped staging (host pointer, device alias)
int32_t* h_tok_ = nullptr; int32_t* m_tok_ = nullptr; // T
int32_t* h_step_ = nullptr; int32_t* m_step_ = nullptr; // T * kStepCount
int32_t* h_pos_ = nullptr; int32_t* m_pos_ = nullptr; // T * n_head
int32_t* h_commit_ = nullptr; int32_t* m_commit_ = nullptr; // [n_keep, n_keep-1, pos_0 .. pos_{T-1}]
float* h_ple_ = nullptr; float* m_ple_ = nullptr; // T * n_embd
int32_t* h_out_ = nullptr; int32_t* m_out_ = nullptr; // T argmax ids
float* h_x_ = nullptr; float* m_x_ = nullptr; // doorbell payload: T * n_embd
int32_t* h_ids_ = nullptr; int32_t* m_ids_ = nullptr; // T * k
float* h_w_ = nullptr; float* m_w_ = nullptr; // T * k
uint32_t* h_seq_ = nullptr; uint32_t* m_seq_ = nullptr;
uint32_t* h_flag_ = nullptr; uint32_t* m_flag_ = nullptr;
uint32_t* h_flagA_ = nullptr; uint32_t* m_flagA_ = nullptr; // the GPU plan is in place
uint32_t* h_flagB_ = nullptr; uint32_t* m_flagB_ = nullptr; // the PCIe share's DMA copies have landed
cudaEvent_t commit_done_ = nullptr; // recorded after an async commit (set_commit_async); see wait_commit
bool commit_pending_ = false;
cudaStream_t copy_ = nullptr; // the copy engine's stream (DMA of missed experts)
struct FlagSet { uint32_t* flag; uint32_t value; };
FlagSet flag_sets_[2 * 64 * 2] = {}; // host-function arguments, one per (layer, group)
static void fetch_dma(void* ctx, const uint8_t* const* src, int n, size_t bytes);
static void raise_flag(uint32_t* flag, uint32_t value);
int32_t* h_plan_ = nullptr; int32_t* m_plan_ = nullptr; // counts | start | dst | tok | ptr (as int32 pairs)
int64_t plan_i32_ = 0; // int32 words in the plan block
GpuPlanSink sink_;
uint32_t cur_layer_ = 0;
static void publish_plan(void* ctx);
void set_plan_slot(int grp);
bool split_ = false; // opt-in (--spec-split): exact but slower, see the overlap study
int groups_[9] = {};
float* h_ymiss_ = nullptr; float* m_ymiss_ = nullptr; // T * k * n_embd
// device
void* arena_ = nullptr;
int32_t *tok_ = nullptr, *step_ = nullptr, *pos_ = nullptr, *commit_ = nullptr;
float *ple_ = nullptr, *emb_ = nullptr, *R_ = nullptr, *mixed_ = nullptr, *bo_ = nullptr;
float *inj_ = nullptr, *inj2_ = nullptr, *lo_ = nullptr, *rs_ = nullptr, *xn_ = nullptr;
uint8_t* xq_ = nullptr; // T columns of q8_1
uint8_t* sh_xq_ = nullptr; // T columns of q8_1 for shared expert branch
float *qkv_L_ = nullptr, *h_L_ = nullptr, *gate_L_ = nullptr, *beta_L_ = nullptr; // per GDN layer
float *z_ = nullptr, *y_ = nullptr, *y_dummy_ = nullptr;
float *qfull_ = nullptr, *qcur_ = nullptr, *kcur_ = nullptr, *vcur_ = nullptr, *idx_raw_L_ = nullptr;
float *qidx_ = nullptr, *scores_ = nullptr, *attn_ = nullptr, *attn32_ = nullptr, *attn_scratch_ = nullptr;
float* tail_snap_ = nullptr; // per QSA layer
int32_t* sel_ = nullptr;
float *logits_ = nullptr, *w_ = nullptr, *shared_ = nullptr, *parts_ = nullptr, *hit_out_ = nullptr;
int32_t *ids_ = nullptr, *hit_slot_ = nullptr, *hit_dst_ = nullptr, *hit_count_ = nullptr;
int32_t* plan_ = nullptr; // device copy of the plan block
uint8_t* staging_ = nullptr; // VRAM slots for the PCIe share of the misses
static constexpr int64_t kStagingBlobs = 16;
static constexpr int64_t kPcieGroupRows = 4; // the PCIe call's groups side by side (of <= 16)
uint8_t* hit_xq_ = nullptr;
uint8_t* nat_xq_ = nullptr; // plan v0.3 P6: q8_1 activations for a native pack's grouped experts
float* hit_xs_ = nullptr;
void* hit_scratch_ = nullptr;
float *head_mixed_ = nullptr, *head_inj_ = nullptr, *head_logits_ = nullptr;
uint16_t* sh_bf16_ = nullptr;
float *sh_gate_ = nullptr, *sh_up_ = nullptr, *sh_g_ = nullptr;
float* hist_snap_ = nullptr; // T * NG_HIST * NG_HC_DIM
int64_t cap_ = 0, max_blocks_ = 0, attn_scratch_floats_ = 0;
};
} // namespace strata::core
// include/strata/core/verify.hpp - plan v0.3 P6: the speculative VERIFY window.
//
// T tokens at consecutive positions p0 .. p0+T-1 - the last accepted token and T-1 drafts - go through all 48
// layers in ONE captured graph, and the head's argmax is produced for every one of them. Token t's argmax is
// what plain greedy decode would produce after token t, BIT FOR BIT: every kernel here is either the single-token
// kernel applied per token, or a multi-token kernel whose per-token arithmetic is the single-token kernel's
// (multi-column MMVQ in exact mode, the T-token GDN kernels, the per-token hit activation, the multi-token CPU
// expert rows). So a draft is accepted exactly when greedy decode would have produced it.
//
// What the window costs is the dense weights read ONCE for T tokens and the union of the T tokens' missed
// experts on the CPU (measured on decode traces: 1.75x one token's misses for T=2, 2.4x for 3, 3.05x for 4).
//
// STATE. The window appends K/V and indexer keys for all T positions and leaves the GDN state untouched.
// `commit(n_keep)` then makes the first `n_keep` tokens permanent: the GDN conv history and recurrent state are
// advanced by replaying those tokens from inputs the window stored, the indexer's key tail is restored from a
// snapshot and the accepted keys re-appended (a rejected key can land in a slot the current block still needs),
// and the PLE history is set to its snapshot after token n_keep-1. K/V cells past the accepted prefix are simply
// overwritten when those positions are processed again, before any query can read them.
//
// Requires the default native decode configuration (native projections, fused GR, fused GDN, fast attention and
// selection, native indexer) and a profile-filled VRAM expert tier with its residency table on the device.
#pragma once
#include <cstdio>
#include "strata/core/expert_source.hpp"
#include "strata/core/layer.hpp"
#include "strata/core/session.hpp"
#include "strata/kernels/sampler.hpp"
#include <cuda_runtime.h>
#include <atomic>
#include <chrono>
#include <map>
#include <cstdint>
#include <string>
#include <vector>
namespace strata::core {
class NativeHead;
class RemoteExpertOpt;
/// The CPU pool for a window: x_f (n_tok, n_embd), ids (n_tok, k) -> out (n_tok * k, n_embd), hit rows zeroed.
using PoolMultiFn = void (*)(void* user, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k, float* out,
int64_t layer);
struct VerifyHits {
const int32_t* d_res = nullptr; ///< device [n_layers * n_expert] slot or -1
const int32_t* h_res = nullptr; ///< host [n_layers * n_expert] slot or -1 (when all >= 0 in stage: zero-doorbell)
const uint8_t* cache_base = nullptr; ///< slot 0 of the VRAM expert arena
const uint64_t* slot_off = nullptr; ///< E-6: host per-slot offsets when slots differ in size (null: slot * blob)
int64_t n_slots = 0; ///< E-6: how many (for the device copy)
int64_t blob = 0;
};
class Verifier {
public:
Verifier() = default;
~Verifier();
Verifier(const Verifier&) = delete;
Verifier& operator=(const Verifier&) = delete;
void set_remote_expert_opt(RemoteExpertOpt* opt) { remote_opt_ = opt; } // before init/capture
/// The watchdog's view of the window in flight (issue #31): the layer, the GPU's sequence, the flags.
void diag(std::FILE* f) const;
/// #649 (STRATA_VERIFY_TRACE=1, opt-in): the last windows' host-side handshake events and the GPU's breadcrumbs
/// (in mapped memory, so they read while the GPU hangs). Nothing when the trace is off.
void trace_dump(std::FILE* f) const;
/// #267: raise every flag the window's spin kernels wait on past any ring (UINT32_MAX), so a window the GPU
/// cannot finish drains instead of staying resident, then wait up to `timeout_ms` for its streams. For the
/// paths that give up on the engine (a timed-out window, the serve watchdog): the window then ran on whatever
/// the flags guarded, so this verifier refuses every later window. True when the streams finished.
bool release_gpu_waits(int timeout_ms);
/// `max_t` <= kVerifyMaxT. `head` may be null (the canonical head is then run per token).
bool init(const WeightTable& wt, const ModelGeometry& g, SessionState& ss, const VerifyHits& hits,
const NativeHead* head, int max_t, std::string& err);
/// One window: `tokens[0..T)` at positions pos0.., the pool served per layer; `out[t]` = argmax after token t.
/// The PLE rows are gathered here from `ss.ple_prev` and the tokens. Captures the T-token graph on first use.
bool run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, void* user, int32_t* out, std::string& err);
/// Diagnostics: row `t` of the last window's head logits (n_vocab floats) to the host. Valid after run().
bool copy_logits(int t, float* host) const;
int64_t vocab() const { return next_ ? next_->vocab() : n_vocab_; }
/// The sampling the verify window's head applies (temperature / top_p / top_k / seed). Set per
/// request; greedy by default. The sampling itself runs OUTSIDE the captured graph - its
/// parameters would otherwise be baked forever - so this can change between requests freely.
void set_sampling(const strata::kernels::SamplerParams& sp) {
sampling_ = sp; // row t of a window at pos0 draws Philox(seed, pos0 + t): see run()
if (next_) next_->set_sampling(sp);
}
/// The penalty histories for `sampling_.penalty_last_n`: ONE ROW PER WINDOW ROW, T rows of `history_len`
/// int32 slots at that stride (`strata::kernels::penalty_rows` builds them), most recent token LAST, unused
/// front slots -1 (the kernel reads only the tail window). Row t follows the window's drafts 1..t - staging
/// row 0 alone (before 0.1.19) left the drafted rows with unwritten histories. Null disables the penalties
/// entirely - the neutral run's sampling call is byte-for-byte what it was. The engine re-uploads the rows
/// before every window; the buffer must hold kVerifyMaxT rows and stay alive across the request.
void set_history(const int32_t* history, int history_len) {
hist_d_ = history;
hist_len_ = history_len;
if (next_) next_->set_history(history, history_len);
}
/// Off: `run` skips the request's head sampling and `out` is the recorded greedy pick. For windows whose
/// picks are discarded - a prompt read through windows commits every token - so they cost no sampler launch
/// or sync and never read a history staged for another position.
void set_head_sampling(bool on) { head_sampling_ = on; if (next_) next_->set_head_sampling(on); }
/// LAYER SPLIT (multi-GPU): this verifier runs layers [layer_begin, layer_end) of every window. A stage that
/// does not start at layer 0 takes its residual from `handoff_in` instead of embedding the tokens; a stage that
/// does not end at the last layer writes its residual to `handoff_out` and has no head. The hand-off holds,
/// per token, the residual R (hc x n_embd), the last layer's pending write bo (n_embd) and inject (hc): the next
/// stage folds that write into its first read exactly as the unsplit window does, so the split is bit-exact.
/// Both pointers must be device-visible (mapped pinned memory, portable when the stages are on two devices).
/// Set before `init`. Default: the whole model, no hand-off.
void set_stage(int64_t layer_begin, int64_t layer_end, const float* handoff_in, float* handoff_out) {
lb_ = layer_begin; le_ = layer_end; hand_in_ = handoff_in; hand_out_ = handoff_out;
}
/// The next stage: `run` and `commit` continue into it (its pool calls get `next_user`); sampling settings
/// and `final_R` are the last stage's.
void set_next(Verifier* next, void* next_user) { next_ = next; next_user_ = next_user; }
/// floats per token in a hand-off buffer
static int64_t handoff_floats(const ModelGeometry& g) { return (int64_t) g.hc * g.n_embd + g.n_embd + g.hc; }
/// Keep the first `n_keep` (1..T) tokens of the last window; advances `ss.ple_prev` by them.
bool commit(int n_keep, std::string& err);
// ================================ SEVERAL SEQUENCES IN ONE WINDOW ================================
//
// A batch window holds S INDEPENDENT sequences, one token each: row s is slot s, at slot s's own position,
// reading and writing slot s's own state (GDN recurrence and conv history, QSA K/V and indexer, PLE history),
// which lives in `slots[s]` - a session carved like this verifier's own (same layer range, same max_cells).
// Everything that is per row already (hyper-connections, dense projections, router, shared expert, the
// routed experts on the GPU and the CPU, the head) runs once over the S rows, so the weights are read once
// per window for all the sequences. Row s's arithmetic is the single-token window's, so a slot's greedy
// tokens are its solo greedy tokens (modulo the multi-token CPU kernel choice: STRATA_IQ_MT_MIN=1).
//
// Greedy only, no drafts (MTP) in a batch window. `init_slots` once after `init` (S <= max_t); a layer
// split's stages each get their own sessions, and run_slots/commit_slots continue into the next stage.
bool init_slots(const std::vector<SessionState*>& slots, std::string& err);
int n_slots() const { return (int) slots_.size(); }
/// One batch window over slots [0, S): tokens[s] at positions pos[s]; out[s] = the greedy pick after it.
bool run_slots(int S, const int32_t* tokens, const int64_t* pos, PoolMultiFn pool, void* user, int32_t* out,
std::string& err);
/// The same over the S slots `rows` (row t is slot rows[t], any distinct slots in any order): the slots not
/// listed are not touched, so an idle slot keeps its state (a finished conversation it may continue later).
bool run_slot_rows(const int* rows, int S, const int32_t* tokens, const int64_t* pos, PoolMultiFn pool, void* user,
int32_t* out, std::string& err);
/// Keep every row of the last batch window: each slot's state advances by its one token.
bool commit_slots(std::string& err);
// ---- The stages of a layer split as a PIPELINE. A batch window over the slot GROUP
// [base, base + S) is launched on ONE stage with its commit right behind it on the stage's stream (a batch window
// keeps every row, so the commit needs no host decision), and the host serves the rings of every stage that has
// a window in flight from one thread (batch_poll does not block). Stage k can then run group g while stage k+1
// runs group g-1. Rows of group `base` use hand-off rows [base, base + S), so groups never share a hand-off row.
bool batch_launch(int base, int S, const int32_t* tokens, const int64_t* pos, std::string& err);
/// 1 = this stage's window and commit are done (the last stage's picks are in batch_out), 0 = still running,
/// -1 = an error (err). Serves every layer that has rung so far.
int batch_poll(PoolMultiFn pool, void* user, std::string& err);
bool batch_busy() const { return b_running_; }
/// A slot's sampling (temperature / top_p / top_k / min_p / seed; penalties are not applied in batch windows):
/// its row is drawn again on the last stage with Philox(seed, position), as a solo window draws it. Greedy by
/// default. Set on the first stage, it reaches the last.
void set_slot_sampling(int slot, const strata::kernels::SamplerParams& sp) {
if (slot >= 0 && slot < (int) slot_sp_.size()) slot_sp_[(size_t) slot] = sp;
if (next_) next_->set_slot_sampling(slot, sp);
}
const int32_t* batch_out() const { return b_out_; }
bool last_stage() const { return g_ != nullptr && le_ == g_->n_layers; }
/// commit() returns without waiting for its graph (a single-GPU session sets it): the next window follows it on
/// the same stream and the drafter reads nothing it writes, so it overlaps the draft. Whoever reads or writes
/// the session from another stream or the host afterwards (a new request, a checkpoint, a snapshot, the prompt
/// path, the end of a run) calls wait_commit() first. STRATA_COMMIT_SYNC=1 keeps the wait.
static void set_commit_async(bool on);
/// Waits for the last commit graph when commit() did not (an event recorded after it, not the whole device);
/// false with `err` when it failed. Free when nothing is pending.
bool wait_commit(std::string& err);
/// Measurement hook (STRATA_LOGPOS): after run(), write one line per row t of the last window's head -
/// "pos target logprob top top_logprob hit extra_logprob target_logprob_without_extra" - where row t is the
/// distribution at pos0 + t, targets[t] is the token at pos0 + t + 1, extra_logprob is the log-probability of
/// `extra_id` in the same row, and the last column is the target's log-probability in that row renormalized
/// over every token but `extra_id` (both nan when they do not apply). A layer split's earlier stage forwards
/// to the stage that holds the head.
bool window_logprobs(const int32_t* targets, int T, int64_t pos0, int32_t extra_id, std::FILE* out,
std::string& err);
/// Token t's residual after the last layer, (hc, n_embd) on the device, valid until the next `run`.
const float* final_R(int t) const;
const float* final_R_all() const { return next_ ? next_->final_R_all() : R_; }
/// The GPU plan the pool writes each layer (VRAM hits + the PCIe share of the misses); give it to the
/// dispatch (`ExpertDispatch::plan`) before the first `run`.
GpuPlanSink* plan_sink() { return &sink_; }
/// Plan v0.3 P6: split the window into two token groups and pipeline the CPU experts of one with the GPU work
/// of the other (default on). Set before the first `run`.
void set_split(bool on) { split_ = on; }
/// Plan v0.3 P6: how the PCIe share of the misses reaches the GPU: 0 = DMA into staging (the copy engine works
/// beside the CPU; best when the CPU is compute-bound, the i-quants), 1 = the grouped kernel reads the mapped
/// arena directly, 2 = a copy kernel stages it inside the graph (no API calls on the pool's thread; best when
/// the CPU is RAM-bound, Q2_0). Set before the first `run`.
void set_pcie_mode(int mode) { sink_.pcie_mode = mode; }
/// the pool never plans a PCIe share (--pcie-frac 0): the window skips that path. Before the first run.
double ms_wait = 0, ms_pool = 0, ms_host = 0, ms_commit = 0;
int64_t windows = 0;
/// STRATA_VERIFY_PROFILE=1 - GPU stage times of the windows since the last call (ms per
/// window), as one line; empty when off.
std::string profile_report();
private:
RemoteExpertOpt* remote_opt_ = nullptr;
bool capture(int T, std::string& err);
// batch windows (see init_slots)
std::vector<SessionState*> slots_;
bool batch_rec_ = false; ///< record_window is capturing a batch window
int row_base_ = 0; ///< ... its hand-off rows start here (a pipeline group's own rows)
int brow_[8] = {}; ///< ... and row t is slot brow_[t]
bool last_batch_ = false; ///< the last run was a batch window (set_plan_slot: one group)
std::map<uint64_t, cudaGraphExec_t> exec_bm_, commit_bm_; ///< key: batch_key(rows, S, hand-off base)
int last_rows_[8] = {}; ///< the slots of the last batch window's rows
static uint64_t batch_key(const int* rows, int S, int hbase) {
uint64_t k = (uint64_t) hbase << 40 | (uint64_t) S << 32;
for (int t = 0; t < S; ++t) k |= (uint64_t) (rows[t] & 15) << (4 * t);
return k;
}
// batch_launch / batch_poll
bool b_running_ = false;
int64_t b_k_ = 0, b_steps_ = 0;
std::chrono::steady_clock::time_point b_last_;
int32_t b_out_[8] = {};
std::vector<strata::kernels::SamplerParams> slot_sp_;
bool sample_rows(int S, std::string& err); ///< the sampled slots' rows of the last batch window
int32_t* h_commitb_ = nullptr; int32_t* m_commitb_ = nullptr; // per slot [1, 0, pos, -1 ..], stride 2 + max_t
int32_t* commitb_ = nullptr;
float* tail_snap_b_ = nullptr; ///< per (slot, QSA layer) indexer tail snapshot
void* arena_b_ = nullptr;
int64_t last_pos_b_[8] = {};
bool capture_batch(const int* rows, int S, int hbase, std::string& err);
void collect_profile(); ///< STRATA_VERIFY_PROFILE: add the last window's stamps to prof_sum_
bool capture_commit_batch(const int* rows, int S, int hbase, std::string& err);
bool stage_batch(const int* rows, int S, int hbase, const int32_t* tokens, const int64_t* pos, std::string& err);
strata::kernels::SamplerParams sampling_ = [] {
strata::kernels::SamplerParams s;
s.greedy = true;
s.temperature = 0.0f;
return s;
}(); ///< greedy by default; per-request via set_sampling
const int32_t* hist_d_ = nullptr; ///< penalty-history row (set_history); null = no penalties apply
int hist_len_ = 0;
bool head_sampling_ = true; ///< set_head_sampling
int device_ = -1; ///< the device `init` ran on: run/commit switch to it (layer split)
std::atomic<bool> released_{false}; ///< #267: release_gpu_waits ran (maybe on the watchdog thread): no more windows
bool all_resident_ = false; ///< 100% of experts in [lb_, le_) resident in VRAM: zero-doorbell graph
bool device_plan_ = false; ///< E-6: resident-only layers planned on the device (STRATA_VERIFY_DEVICE_PLAN)
uint32_t* skip_ = nullptr; ///< E-6: per group, the ring whose plan the device built (0: the host's)
unsigned long long* slot_off_d_ = nullptr; ///< E-6: the slot offsets on the device
int64_t lb_ = 0, le_ = -1; ///< set_stage: the layers this verifier runs (-1: to the last)
const float* hand_in_ = nullptr;
float* hand_out_ = nullptr;
Verifier* next_ = nullptr;
void* next_user_ = nullptr;
bool ple_stage() const { return lb_ <= 1 && 1 < le_; } ///< holds layer 1, where the PLE block runs
void stage_inputs(int T, const int32_t* tokens, int64_t pos0);
bool staged_ = false;
bool prelaunched_ = false;
bool copy_used_ = false;
bool capture_commit(std::string& err);
bool record_window(int T, cudaStream_t cs, std::string& err);
// #649: STRATA_VERIFY_TRACE=1 - a host event ring (trace_ev) and GPU breadcrumbs: the profiler's stamp points,
// per layer and token group, written to mapped memory (null when the trace is off)
void trace_ev(const char* what, int64_t step, int64_t layer, int64_t aux) const;
unsigned long long* trace_h_ = nullptr;
unsigned long long* trace_m_ = nullptr;
size_t trace_n_ = 0;
static constexpr int kProfPer = 33; // stamps per layer (32 left the hc-read second
// half's up-stamp at slot 32 = the next layer's slot 0: D8)
bool prof_on_ = false;
unsigned long long* prof_ = nullptr; // device: n_layers * kProfPer + 4 stamps
std::vector<unsigned long long> prof_h_;
double prof_sum_[2][kProfPer] = {}; // [GDN / QSA layers][stage]
int64_t prof_windows_ = 0;
const WeightTable* wt_ = nullptr;
const ModelGeometry* g_ = nullptr;
SessionState* ss_ = nullptr;
VerifyHits hits_;
const NativeHead* head_ = nullptr;
int max_t_ = 0;
int last_t_ = 0;
int64_t last_pos0_ = 0;
int32_t last_tokens_[8] = {};
int64_t n_vocab_ = 0;
cudaStream_t cs_ = nullptr;
cudaStream_t sh_cs_ = nullptr;
cudaEvent_t ev_fork_ = nullptr, ev_join_ = nullptr, handoff_ev_ = nullptr;
cudaGraphExec_t exec_[9] = {};
cudaGraphExec_t commit_exec_ = nullptr;
// mapped staging (host pointer, device alias)
int32_t* h_tok_ = nullptr; int32_t* m_tok_ = nullptr; // T
int32_t* h_step_ = nullptr; int32_t* m_step_ = nullptr; // T * kStepCount
int32_t* h_pos_ = nullptr; int32_t* m_pos_ = nullptr; // T * n_head
int32_t* h_commit_ = nullptr; int32_t* m_commit_ = nullptr; // [n_keep, n_keep-1, pos_0 .. pos_{T-1}]
float* h_ple_ = nullptr; float* m_ple_ = nullptr; // T * n_embd
int32_t* h_out_ = nullptr; int32_t* m_out_ = nullptr; // T argmax ids
float* h_x_ = nullptr; float* m_x_ = nullptr; // doorbell payload: T * n_embd
int32_t* h_ids_ = nullptr; int32_t* m_ids_ = nullptr; // T * k
float* h_w_ = nullptr; float* m_w_ = nullptr; // T * k
uint32_t* h_seq_ = nullptr; uint32_t* m_seq_ = nullptr;
uint32_t* h_flag_ = nullptr; uint32_t* m_flag_ = nullptr;
uint32_t* h_flagA_ = nullptr; uint32_t* m_flagA_ = nullptr; // the GPU plan is in place
uint32_t* h_flagB_ = nullptr; uint32_t* m_flagB_ = nullptr; // the PCIe share's DMA copies have landed
cudaEvent_t commit_done_ = nullptr; // recorded after an async commit (set_commit_async); see wait_commit
bool commit_pending_ = false;
cudaStream_t copy_ = nullptr; // the copy engine's stream (DMA of missed experts)
struct FlagSet { uint32_t* flag; uint32_t value; };
FlagSet flag_sets_[2 * 64 * 2] = {}; // host-function arguments, one per (layer, group)
static void fetch_dma(void* ctx, const uint8_t* const* src, int n, size_t bytes);
static void raise_flag(uint32_t* flag, uint32_t value);
int32_t* h_plan_ = nullptr; int32_t* m_plan_ = nullptr; // counts | start | dst | tok | ptr (as int32 pairs)
int64_t plan_i32_ = 0; // int32 words in the plan block
GpuPlanSink sink_;
uint32_t cur_layer_ = 0;
static void publish_plan(void* ctx);
void set_plan_slot(int grp);
bool split_ = false; // opt-in (--spec-split): exact but slower, see the overlap study
int groups_[9] = {};
float* h_ymiss_ = nullptr; float* m_ymiss_ = nullptr; // T * k * n_embd
// device
void* arena_ = nullptr;
int32_t *tok_ = nullptr, *step_ = nullptr, *pos_ = nullptr, *commit_ = nullptr;
float *ple_ = nullptr, *emb_ = nullptr, *R_ = nullptr, *mixed_ = nullptr, *bo_ = nullptr;
float *inj_ = nullptr, *inj2_ = nullptr, *lo_ = nullptr, *rs_ = nullptr, *xn_ = nullptr;
uint8_t* xq_ = nullptr; // T columns of q8_1
uint8_t* sh_xq_ = nullptr; // T columns of q8_1 for shared expert branch
float *qkv_L_ = nullptr, *h_L_ = nullptr, *gate_L_ = nullptr, *beta_L_ = nullptr; // per GDN layer
float *z_ = nullptr, *y_ = nullptr, *y_dummy_ = nullptr;
float *qfull_ = nullptr, *qcur_ = nullptr, *kcur_ = nullptr, *vcur_ = nullptr, *idx_raw_L_ = nullptr;
float *qidx_ = nullptr, *scores_ = nullptr, *attn_ = nullptr, *attn32_ = nullptr, *attn_scratch_ = nullptr;
float* tail_snap_ = nullptr; // per QSA layer
int32_t* sel_ = nullptr;
float *logits_ = nullptr, *w_ = nullptr, *shared_ = nullptr, *parts_ = nullptr, *hit_out_ = nullptr;
int32_t *ids_ = nullptr, *hit_slot_ = nullptr, *hit_dst_ = nullptr, *hit_count_ = nullptr;
int32_t* plan_ = nullptr; // device copy of the plan block
uint8_t* staging_ = nullptr; // VRAM slots for the PCIe share of the misses
static constexpr int64_t kStagingBlobs = 16;
static constexpr int64_t kPcieGroupRows = 4; // the PCIe call's groups side by side (of <= 16)
uint8_t* hit_xq_ = nullptr;
uint8_t* nat_xq_ = nullptr; // plan v0.3 P6: q8_1 activations for a native pack's grouped experts
float* hit_xs_ = nullptr;
void* hit_scratch_ = nullptr;
float *head_mixed_ = nullptr, *head_inj_ = nullptr, *head_logits_ = nullptr;
uint16_t* sh_bf16_ = nullptr;
float *sh_gate_ = nullptr, *sh_up_ = nullptr, *sh_g_ = nullptr;
float* hist_snap_ = nullptr; // T * NG_HIST * NG_HC_DIM
int64_t cap_ = 0, max_blocks_ = 0, attn_scratch_floats_ = 0;
};
} // namespace strata::core
src/core/mtp.cpp
+1030 -993
View File
File diff suppressed because it is too large. Load diff
src/core/verify.cpp
+2266 -2158
View File
File diff suppressed because it is too large. Load diff