mirror of
https://github.com/Helldez/BigMoeOnEdge.git
synced 2026-10-03 11:35:50 +00:00
45a90a2df5
* feat(moe): --drop-cold-experts, spend quality only where it buys I/O Turbo top-k drops the tail of a routing whether or not those experts were already in RAM. A resident expert costs no flash read, so that trade pays quality for nothing on the ~80% of decode routings that are cache hits. This adds the cache-aware version: skip a routed expert only when it is a cache MISS and the router weighted it below frac x (1/top-k). Replayed over the committed route traces at frac 1.0, decode phase, that avoids 66% of flash reads for 9.5% of the router's weight mass, against 59%/37% for --n-expert-used 3 — about 3x the reads avoided at a comparable cost. The threshold is a curve, not a switch: 0.75 trades 4.4% of the mass for 37% of the reads, better than --n-expert-used 5 on both axes. Implementation. The decision needs the FINAL router weights, which arrive several nodes after the topk where the streamer normally loads, so with the policy armed load_layer() is deferred to the terminal node of the layer's weight chain. Which node that is depends on the model's gating, so the hook learns it from the graph rather than carrying an architecture table; until it is known a layer loads at its topk node undropped. A dropped slot has its weight zeroed and its expert id repointed at the routing's top-weighted expert: an expert we decline to read may sit in reserved-but-uncommitted VM and mul_mat_id would still touch it, so the kernel is given memory that is certainly resident and multiplies it by exactly zero. Survivors are rescaled by default, since a systematically shrunk expert output perturbs the residual stream more than the missing contribution does. Prefill is excluded by default (cold cache, ~4x the weight mass discarded, and compute-bound anyway). The largest weight in a routing is always at least the uniform share, so frac <= 1 can never empty a layer; validate() enforces the bound and the top expert is pinned regardless. Gates: G8a proves the deferral and the learned terminal node are transparent (a threshold below any producible weight leaves the output byte-identical), G8b that full strength with the cache off never reaches an unloaded expert. Unlike every other knob this one is state-dependent: what gets dropped depends on what the cache held, so output is not reproducible across runs. Off by default, not in the app's settings, and NOT yet measured on device — the numbers above are a static replay and an upper bound. docs/expert- dropping.md states what is owed before it is recommended anywhere. * feat(app): expose cache-aware expert dropping in Settings Speed / quality -> Drop cold experts, as a percentage of the uniform share (off / 50 / 75 / 100). The engine takes a fraction; the app stores integer rungs, so the setting divides by 100 on the way to the flag. Disabled in mmap mode: the policy asks the expert source what is resident, and there is no expert source without the streamer. Included in the session signature, so changing it reopens the session rather than being ignored by a process already loaded. Off by default. This exists so the A/B can be run where the engine actually ships -- through the app, not a pushed CLI binary. * fix(moe): require the cache for dropping, and correct what it reports Review of the first two commits found the policy could be armed in a configuration where it is not cache-aware at all, and that two of the numbers it reports were wrong. - Require the LRU cache. With --cache-mb 0 query_residency answers all-miss, so the policy silently degenerated into an unconditional weight cut -- exactly what --n-expert-used already does, under a flag claiming to consult residency. validate() now rejects it, as it already did for --prefetch, and the app gates the setting on the same condition. - Fix experts_routed. It was incremented inside apply_drop, so it counted what the policy examined rather than what the router selected: layers before the terminal weight node is learned, and every un-armed phase, were missing from the denominator. The reported drop rate was a fraction of the wrong thing. - Re-learn instead of re-betting. If the node learned as terminal does not arrive, the deferral now also forgets it, so the next graph loads at the topk node while it re-learns. Deferring again on a stale guess would repeat the fault every token against a graph that had moved. - Point the gates at a real cache. G8a/G8b ran with the cache off, where the shared-slot path has no reserved-but-uncommitted memory -- so the id repointing, which is the design's whole safety argument, was never exercised. They now run against a constantly-evicting budget. Adds G8a' (asserts routings were examined and none dropped, so an inert-threshold flake fails legibly) and G8c (at top-k 1 dropping is a proven no-op, pinning both the top-expert guarantee and the threshold being taken against the effective top-k). Docs: three metrics change meaning under dropping and none of them said so. A dropped routing is a miss that is never looked up, so cache_hit_pct rises without the cache serving more, and token/layer_demand measure what was staged rather than routed -- documented in telemetry.md, pressure.md (size the cache with dropping off, then turn it on) and metrics.h. prefetch.md's "cannot change output" is scoped: under dropping a correct guess un-drops an expert. limitations.md gains the non-reproducibility entry, benchmark-method.md the axis plus a warning that reversing the run order cannot distinguish a moved drop rate from a contaminated cell, and architecture.md/runtime.h no longer claim unconditional determinism. Fixes two anchors the README rename broke, and a changelog sentence that quoted the equal-I/O row while drawing the equal-quality conclusion. App: Drop cold experts defaults to 75%. The default is a product decision taken on the maintainer's device; no benchmark for it is published here, and docs/expert-dropping.md says that plainly instead of implying a measured figure. The CLI stays off by default -- the byte-identity gates need a deterministic default.
980 lines
47 KiB
C++
980 lines
47 KiB
C++
#include "bmoe/session.h"
|
|
#include "bmoe/recipe.h"
|
|
#include "bmoe/route_trace.h"
|
|
#include "bmoe/decode_trace.h"
|
|
#include "chat_parse.h"
|
|
#include "thinking_control.h"
|
|
#include "../moe/router_hook.h"
|
|
#include "../moe/expert_stream_source.h"
|
|
#include "../moe/gguf_offsets.h"
|
|
#include "../io/platform_io.h"
|
|
|
|
#include "llama.h"
|
|
#include "ggml.h"
|
|
|
|
// llama.cpp's `common` layer (NOT the stable public API): chat-template rendering and
|
|
// reasoning parsing. See the note in the root CMakeLists / docs/seam.md.
|
|
#include "chat.h"
|
|
#include "common.h"
|
|
|
|
#include <algorithm>
|
|
#include <atomic>
|
|
#include <chrono>
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <exception>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <unordered_set>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
namespace bmoe {
|
|
|
|
namespace {
|
|
|
|
using clock_t_ = std::chrono::steady_clock;
|
|
double secs(clock_t_::time_point a, clock_t_::time_point b) {
|
|
return std::chrono::duration<double>(b - a).count();
|
|
}
|
|
|
|
llama_token argmax(const float * logits, int n_vocab) {
|
|
llama_token best = 0;
|
|
float best_v = logits[0];
|
|
for (int v = 1; v < n_vocab; ++v)
|
|
if (logits[v] > best_v) {
|
|
best_v = logits[v];
|
|
best = v;
|
|
}
|
|
return best;
|
|
}
|
|
|
|
// The generation phase's running measurement state, and the one place a generated token's cost is
|
|
// written down. The cursors hold the source's absolute totals as of the previous token — its stats
|
|
// are cumulative across a warm session, so every per-token flash figure is a delta against these —
|
|
// and the totals are what the summary averages over n_gen. Grouped into one object so the per-token
|
|
// metrics block can be a method instead of ten more locals threaded through generate().
|
|
struct GenTally {
|
|
// Fixed for the run; kept here so record() needs only the token's own measurements.
|
|
bool overlap = false;
|
|
int n_threads = 1;
|
|
|
|
long long prev_bytes = 0;
|
|
double prev_io_s = 0.0;
|
|
double prev_mgmt_s = 0.0;
|
|
double prev_stall_s = 0.0;
|
|
|
|
uint64_t read_bytes = 0;
|
|
double io_seconds = 0.0;
|
|
double mgmt_seconds = 0.0;
|
|
double stall_seconds = 0.0;
|
|
uint64_t majflt = 0;
|
|
double cpu_seconds = 0.0;
|
|
|
|
// Fill in everything a generated token is measured by — its wall/fault/CPU decomposition, the
|
|
// memory picture, and (when streaming) the flash figures taken as deltas against the cursors
|
|
// above — then advance those cursors and the run totals. The token's TEXT stays with the
|
|
// caller: what a token says depends on chat state, what it cost does not.
|
|
//
|
|
// `wall` is the decode's wall time in seconds and `faults`/`cpu_s` the deltas measured around
|
|
// that same decode; `st` is the expert source's stats, or null when streaming is off.
|
|
void
|
|
record(TokenMetrics & m, double wall, uint64_t faults, double cpu_s, int turn, const IExpertSource::Stats * st) {
|
|
m.wall_ms = wall * 1000.0;
|
|
// Fault/CPU decomposition is independent of streaming — dense-weight faults show up in the
|
|
// mmap baseline too — so record it for every token before the moe/no-moe split below.
|
|
m.majflt = faults;
|
|
m.cpu_ms = cpu_s * 1000.0;
|
|
m.majflt_mib = (double) m.majflt * (double) pio::fault_bytes() / (1024.0 * 1024.0);
|
|
m.turn = turn;
|
|
majflt += m.majflt;
|
|
cpu_seconds += cpu_s;
|
|
|
|
// Read the memory picture AFTER the decode, outside the caller's timing bracket: two /proc
|
|
// reads are cheap but they are not this token's work, and billing them to wall_ms would
|
|
// corrupt the very number the reader is here to trust.
|
|
pio::ProcessMemory pm;
|
|
if (pio::process_memory(&pm)) {
|
|
m.rss_mib = pm.rss_bytes / (1024.0 * 1024.0);
|
|
m.rss_anon_mib = pm.rss_anon_bytes / (1024.0 * 1024.0);
|
|
m.rss_file_mib = pm.rss_file_bytes / (1024.0 * 1024.0);
|
|
m.swap_mib = pm.swap_bytes / (1024.0 * 1024.0);
|
|
}
|
|
pio::DeviceMemory dm;
|
|
if (pio::device_memory(&dm)) {
|
|
m.mem_available_mib = dm.available_bytes / (1024.0 * 1024.0);
|
|
m.mem_free_mib = dm.free_bytes / (1024.0 * 1024.0);
|
|
m.swap_free_mib = dm.swap_free_bytes / (1024.0 * 1024.0);
|
|
}
|
|
if (!st) {
|
|
m.compute_ms = m.wall_ms;
|
|
m.cache_hit_pct = -1.0;
|
|
return;
|
|
}
|
|
|
|
m.dense_resident_frac = st->dense_resident_frac;
|
|
m.cache_budget_mib = st->cache_budget_bytes / (1024.0 * 1024.0);
|
|
m.read_bytes = (uint64_t) ((long long) st->read_bytes - prev_bytes);
|
|
m.io_ms = (st->read_seconds - prev_io_s) * 1000.0;
|
|
m.mgmt_ms = (st->mgmt_seconds - prev_mgmt_s) * 1000.0;
|
|
if (overlap) {
|
|
m.stall_ms = (st->stall_seconds - prev_stall_s) * 1000.0 / n_threads;
|
|
m.compute_ms = m.wall_ms - m.stall_ms - m.mgmt_ms;
|
|
} else {
|
|
m.compute_ms = m.wall_ms - m.io_ms - m.mgmt_ms;
|
|
}
|
|
if (m.compute_ms < 0) m.compute_ms = 0;
|
|
m.cache_hit_pct = st->cache_lookups > 0 ? 100.0 * st->cache_hits / st->cache_lookups : -1.0;
|
|
|
|
prev_bytes = (long long) st->read_bytes;
|
|
prev_io_s = st->read_seconds;
|
|
prev_mgmt_s = st->mgmt_seconds;
|
|
prev_stall_s = st->stall_seconds;
|
|
read_bytes += m.read_bytes;
|
|
io_seconds += m.io_ms / 1000.0;
|
|
mgmt_seconds += m.mgmt_ms / 1000.0;
|
|
stall_seconds += m.stall_ms / 1000.0;
|
|
}
|
|
};
|
|
|
|
// The names of the expert weight tensors the streamer rebinds. "Dense" is defined by subtraction —
|
|
// everything the model has that is NOT one of these — so both consumers below start by asking this
|
|
// same question, and used to answer it with their own copy of the same triple loop.
|
|
std::unordered_set<std::string> expert_tensor_names(const std::vector<LayerExperts> & layers) {
|
|
std::unordered_set<std::string> names;
|
|
for (const LayerExperts & L : layers) {
|
|
if (!L.bound) continue;
|
|
for (int p = 0; p < MoeRecipe::max_exps; ++p)
|
|
if (L.proj[p].tensor) names.insert(L.proj[p].tensor->name);
|
|
}
|
|
return names;
|
|
}
|
|
|
|
// Each layer's bytes that the streamer does NOT manage: everything under blk.<il>. except the
|
|
// expert weight tensors it rebinds — attention, norms, the router, and any per-expert scale left
|
|
// mmap-resident. This is what the layer costs to page in, and it is a static property of the
|
|
// file: nothing about decoding changes it, which is why the route trace states it once in the
|
|
// static block instead of pretending to measure it per step.
|
|
std::vector<uint64_t>
|
|
dense_bytes_per_layer(const GgufOffsets & offs, const std::vector<LayerExperts> & layers, int n_layer) {
|
|
const std::unordered_set<std::string> streamed = expert_tensor_names(layers);
|
|
std::vector<uint64_t> out((size_t) std::max(0, n_layer), 0);
|
|
for (const auto & kv : offs.size_by_name) {
|
|
int il = -1;
|
|
if (std::sscanf(kv.first.c_str(), "blk.%d.", &il) != 1) continue;
|
|
if (il < 0 || il >= n_layer || streamed.count(kv.first)) continue;
|
|
out[(size_t) il] += kv.second;
|
|
}
|
|
return out;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
// All native state lives here, behind the pimpl. Built once by Session::open(); every
|
|
// generate() reuses the loaded model, the live context and the warm expert cache.
|
|
struct Session::Impl {
|
|
SessionConfig cfg;
|
|
std::string arch;
|
|
double load_seconds = 0.0;
|
|
|
|
// Ownership order matters at teardown: the source's I/O pool holds fds into the mmap'd
|
|
// file and its buffers back the rebound expert tensors, so it must be shut down before
|
|
// the context and model are freed. The destructor does that explicitly.
|
|
std::unique_ptr<llama_model, void (*)(llama_model *)> model{nullptr, llama_model_free};
|
|
std::unique_ptr<llama_context, void (*)(llama_context *)> ctx{nullptr, llama_free};
|
|
std::unique_ptr<RouterHook> hook; // heap: its address is baked into cparams.cb_eval_user_data
|
|
ExpertStreamSource source;
|
|
|
|
const llama_vocab * vocab = nullptr;
|
|
int n_vocab = 0;
|
|
int n_layer = 0;
|
|
bool chat_on = false;
|
|
common_chat_templates_ptr chat_tmpls;
|
|
// How a think=false request can be honoured on this model; probed once at open(). Template
|
|
// (the fail-open default) means the flag alone does the job and generate() adds nothing.
|
|
ThinkControl think_ctl = ThinkControl::Template;
|
|
bool backend_inited = false;
|
|
|
|
// Sampling chain, built once at open() only when sampling is requested (temp > 0); null on the
|
|
// greedy default, where the decode loop stays on the argmax fast path. See open()/generate().
|
|
llama_sampler * smpl = nullptr;
|
|
|
|
// Multi-turn chat state (chat mode only). chat_history is the running conversation the
|
|
// template is re-rendered over each turn; kv_tokens mirrors the tokens currently decoded
|
|
// into the context's KV (seq 0), in order, so the next turn can reuse the common prefix
|
|
// and prefill only the diverging suffix instead of re-running the whole conversation.
|
|
std::vector<common_chat_msg> chat_history;
|
|
std::vector<llama_token> kv_tokens;
|
|
|
|
// Route trace (diagnostics): null unless requested AND streaming is on — there is no routing
|
|
// to trace otherwise.
|
|
IRouteTraceSink * route_trace = nullptr;
|
|
|
|
// Which generate() this is, 0-based. It labels a trace's rows, and it labels every token's
|
|
// metrics — a multi-turn CSV is unreadable without it, and the two-turn A/B (a fast turn, an
|
|
// idle, then the turn that pays for it) is exactly what this engine is measured by.
|
|
int turn = 0;
|
|
|
|
// Decode traces (diagnostics): null unless requested. The compute trace needs no streaming —
|
|
// it measures the graph, which a dense mmap run has too; the I/O trace needs the streamer,
|
|
// since there are no engine-issued reads without it. See bmoe/decode_trace.h.
|
|
IComputeTraceSink * compute_trace = nullptr;
|
|
IIoTraceSink * io_trace = nullptr;
|
|
std::vector<IoTraceRow> io_rows_scratch;
|
|
|
|
// Stated once to a metrics sink, before the first token: the model and configuration every row
|
|
// it writes was produced under. info_sent guards a session's many generate() calls from
|
|
// interleaving preambles between turns.
|
|
RunInfo info;
|
|
bool info_sent = false;
|
|
|
|
std::atomic<bool> cancel_requested{false};
|
|
|
|
~Impl() {
|
|
// Deterministic teardown order: stop the I/O pool (it holds fds into the mmap and its
|
|
// buffers back the rebound expert tensors), then the context (its eval callback points
|
|
// at the hook), then the hook, then unmap the model, then release the backend.
|
|
source.shutdown();
|
|
if (smpl) llama_sampler_free(smpl); // independent of ctx/model; free before them
|
|
ctx.reset();
|
|
hook.reset();
|
|
model.reset();
|
|
if (backend_inited) llama_backend_free();
|
|
}
|
|
};
|
|
|
|
Session::Session() : impl_(std::make_unique<Impl>()) {}
|
|
Session::~Session() = default;
|
|
|
|
double Session::load_seconds() const {
|
|
return impl_->load_seconds;
|
|
}
|
|
const std::string & Session::arch() const {
|
|
return impl_->arch;
|
|
}
|
|
int Session::n_ctx() const {
|
|
return impl_->cfg.n_ctx;
|
|
}
|
|
ThinkControl Session::think_control() const {
|
|
return impl_->think_ctl;
|
|
}
|
|
void Session::set_cache_budget_mb(int mib) {
|
|
impl_->source.set_cache_budget((size_t) std::max(0, mib) * 1024ull * 1024ull);
|
|
}
|
|
void Session::cancel() {
|
|
impl_->cancel_requested.store(true, std::memory_order_relaxed);
|
|
}
|
|
|
|
std::unique_ptr<Session> Session::open(const SessionConfig & cfg,
|
|
std::string & error,
|
|
IRouteTraceSink * route_trace,
|
|
IComputeTraceSink * compute_trace,
|
|
IIoTraceSink * io_trace) {
|
|
auto fail = [&](std::string msg) -> std::unique_ptr<Session> {
|
|
error = std::move(msg);
|
|
return nullptr;
|
|
};
|
|
|
|
// Create the session first so its Impl destructor owns backend teardown from this point
|
|
// on: any failure below returns nullptr, destroying `self`, which frees the backend once.
|
|
std::unique_ptr<Session> self(new Session());
|
|
Impl & im = *self->impl_;
|
|
im.cfg = cfg;
|
|
|
|
// llama_backend_init/free is process-global and reference counted; init here, free in ~Impl.
|
|
llama_backend_init();
|
|
im.backend_inited = true;
|
|
|
|
const auto t_load0 = clock_t_::now();
|
|
|
|
// The gguf header answers three separate questions below — the arch-prefixed key for a top-k
|
|
// override, the route trace's effective top-k, and the run info's top-k/expert count — and each
|
|
// used to reopen and reparse the file for its own answer. Read it at most once, lazily: the
|
|
// callers are conditional (a run with no override and no trace asks nothing), so an eager read
|
|
// would be work the common path never needs.
|
|
GgufModelInfo gguf_info;
|
|
bool gguf_info_read = false;
|
|
auto gguf = [&]() -> const GgufModelInfo & {
|
|
if (!gguf_info_read) {
|
|
gguf_info = read_gguf_model_info(cfg.model_path.c_str());
|
|
gguf_info_read = true;
|
|
}
|
|
return gguf_info;
|
|
};
|
|
|
|
// Load with the layout the streamer requires: file-backed mmap, no repack (a repacked
|
|
// q4_K buffer would break the rebind), experts on CPU.
|
|
llama_model_params mparams = llama_model_default_params();
|
|
mparams.use_mmap = true;
|
|
mparams.use_extra_bufts = false;
|
|
mparams.n_gpu_layers = 0;
|
|
|
|
// Optional active-expert override: reduce the model's top-k routing (e.g. 8 -> 6) to cut
|
|
// per-token compute and — under streaming — flash I/O, at a quality cost. Applied purely
|
|
// through llama.cpp's public kv_overrides on the arch-prefixed expert_used_count key: the
|
|
// graph then routes to fewer experts and the whole streaming path adapts automatically
|
|
// (the router hook reads the top-k width from the graph). The array must outlive the load
|
|
// call below. See docs/adding-a-model.md — no llama.cpp patch, no per-arch constants.
|
|
llama_model_kv_override kv_overrides[2];
|
|
std::memset(kv_overrides, 0, sizeof(kv_overrides)); // second entry stays the key[0]==0 terminator
|
|
if (cfg.n_expert_used > 0) {
|
|
const GgufModelInfo & info = gguf();
|
|
if (!info.ok) return fail("cannot read gguf metadata: " + cfg.model_path);
|
|
if (info.arch.empty()) return fail("gguf has no general.architecture; cannot set n_expert_used");
|
|
if (info.n_expert <= 0)
|
|
return fail("n_expert_used was set but the model is not MoE (no " + info.arch + ".expert_count)");
|
|
if (cfg.n_expert_used > info.n_expert)
|
|
return fail("n_expert_used=" + std::to_string(cfg.n_expert_used) + " exceeds the model's expert count (" +
|
|
std::to_string(info.n_expert) + ")");
|
|
const std::string key = info.arch + ".expert_used_count";
|
|
kv_overrides[0].tag = LLAMA_KV_OVERRIDE_TYPE_INT;
|
|
std::snprintf(kv_overrides[0].key, sizeof(kv_overrides[0].key), "%s", key.c_str());
|
|
kv_overrides[0].val_i64 = cfg.n_expert_used;
|
|
mparams.kv_overrides = kv_overrides;
|
|
}
|
|
|
|
llama_model * model = llama_model_load_from_file(cfg.model_path.c_str(), mparams);
|
|
if (!model) return fail("failed to load model: " + cfg.model_path);
|
|
im.model.reset(model);
|
|
|
|
im.vocab = llama_model_get_vocab(model);
|
|
im.n_vocab = llama_vocab_n_tokens(im.vocab);
|
|
im.n_layer = llama_model_n_layer(model);
|
|
|
|
char arch[128] = {0};
|
|
llama_model_meta_val_str(model, "general.architecture", arch, sizeof(arch));
|
|
im.arch = arch;
|
|
|
|
const MoeRecipe * recipe = nullptr;
|
|
if (cfg.moe.enabled) {
|
|
recipe = find_moe_recipe(arch);
|
|
if (!recipe)
|
|
return fail(std::string("no MoE recipe for architecture '") + arch +
|
|
"' — add one in core/src/moe/arch_registry.cpp (see docs/adding-a-model.md)");
|
|
}
|
|
|
|
// Chat templates are model-bound: initialise once here, apply per prompt in generate().
|
|
if (cfg.chatml) {
|
|
try {
|
|
im.chat_tmpls = common_chat_templates_init(model, "");
|
|
im.chat_on = true;
|
|
// Which "thinking off" mechanism this template supports is a property of the model, so
|
|
// it is settled once here rather than re-derived on every turn.
|
|
im.think_ctl = detail::probe_think_control(im.chat_tmpls.get());
|
|
} catch (const std::exception & e) {
|
|
std::fprintf(stderr, "bmoe: chat template unavailable (%s); using raw prompts\n", e.what());
|
|
im.chat_on = false;
|
|
}
|
|
}
|
|
|
|
im.hook = std::make_unique<RouterHook>(recipe ? *recipe : MoeRecipe{}, im.n_layer);
|
|
im.hook->set_prefetch_layers(cfg.moe.prefetch_layers);
|
|
im.hook->set_drop_policy(cfg.moe.drop_cold_frac, cfg.moe.drop_renorm, cfg.moe.drop_prefill);
|
|
|
|
llama_context_params cparams = llama_context_default_params();
|
|
cparams.n_ctx = cfg.n_ctx;
|
|
cparams.n_batch = cfg.n_batch;
|
|
cparams.n_ubatch = cfg.n_batch;
|
|
// The streamer needs the callback to see routing; the compute trace needs it to time nodes.
|
|
// Installing it for the trace alone is what lets a NON-streamed run be measured — the dense
|
|
// mmap baseline the streamed numbers are argued against.
|
|
if (cfg.moe.enabled || compute_trace) {
|
|
cparams.cb_eval = &RouterHook::c_eval;
|
|
cparams.cb_eval_user_data = im.hook.get();
|
|
}
|
|
|
|
llama_context * ctx = llama_init_from_model(model, cparams);
|
|
if (!ctx) return fail("failed to create context");
|
|
im.ctx.reset(ctx);
|
|
llama_set_n_threads(ctx, cfg.n_threads, cfg.n_threads);
|
|
|
|
// Opt-in sampling. temp <= 0 leaves smpl null and the decode loop on argmax — the deterministic
|
|
// default the byte-identity gates rely on. temp > 0 builds the standard chain, using only the
|
|
// public llama_sampler_* API (hard rule 1): common_sampler lives in the non-stable common layer.
|
|
static_assert(SamplingConfig{}.seed == LLAMA_DEFAULT_SEED,
|
|
"SamplingConfig::seed default must mirror LLAMA_DEFAULT_SEED");
|
|
if (cfg.sampling.temp > 0.0f) {
|
|
im.smpl = llama_sampler_chain_init(llama_sampler_chain_default_params());
|
|
llama_sampler_chain_add(im.smpl, llama_sampler_init_top_k(cfg.sampling.top_k));
|
|
llama_sampler_chain_add(im.smpl, llama_sampler_init_top_p(cfg.sampling.top_p, /*min_keep*/ 1));
|
|
llama_sampler_chain_add(im.smpl, llama_sampler_init_temp(cfg.sampling.temp));
|
|
llama_sampler_chain_add(im.smpl, llama_sampler_init_dist(cfg.sampling.seed));
|
|
}
|
|
|
|
// One abort callback for the session's whole life, checking two independent predicates:
|
|
// an explicit cancel() request (any mode) and a fatal streaming I/O error (overlap only).
|
|
// Installing it unconditionally is what lets cancel() interrupt a serial decode too.
|
|
llama_set_abort_callback(
|
|
ctx,
|
|
[](void * ud) -> bool {
|
|
auto * p = static_cast<Impl *>(ud);
|
|
return p->cancel_requested.load(std::memory_order_relaxed) || p->source.fatal();
|
|
},
|
|
&im);
|
|
|
|
if (cfg.moe.enabled) {
|
|
// Capture warm-up: one mmap-resident decode so the eval-callback can harvest the expert
|
|
// tensor pointers from the graph. Any valid token builds the same graph (the expert
|
|
// tensor structure is prompt-independent), so use BOS; KV is wiped afterwards.
|
|
im.hook->begin_capture();
|
|
llama_token warm_tok = llama_vocab_bos(im.vocab);
|
|
if (warm_tok < 0) warm_tok = 0;
|
|
llama_batch warm = llama_batch_get_one(&warm_tok, 1);
|
|
if (llama_decode(ctx, warm) != 0) return fail("capture warm-up decode failed");
|
|
im.hook->end_capture();
|
|
|
|
GgufOffsets offs = read_gguf_offsets(cfg.model_path.c_str());
|
|
if (!offs.ok) return fail("cannot read gguf offsets: " + cfg.model_path);
|
|
|
|
std::vector<LayerExperts> layers = im.hook->captured();
|
|
int n_expert = 0;
|
|
int n_bound = 0;
|
|
for (LayerExperts & L : layers) {
|
|
if (!L.bound) continue;
|
|
++n_bound;
|
|
for (int p = 0; p < MoeRecipe::max_exps; ++p) {
|
|
if (!recipe->exps_suffix[p]) continue; // slot unused by this architecture
|
|
ggml_tensor * t = L.proj[p].tensor;
|
|
if (!t)
|
|
return fail(std::string("captured MoE layer is missing expert tensor '") + recipe->exps_suffix[p] +
|
|
"'");
|
|
auto it = offs.off_by_name.find(t->name);
|
|
if (it == offs.off_by_name.end()) return fail(std::string("no gguf offset for tensor ") + t->name);
|
|
L.proj[p].file_off = it->second;
|
|
const int ne2 = (int) t->ne[2];
|
|
if (n_expert == 0)
|
|
n_expert = ne2;
|
|
else if (ne2 != n_expert)
|
|
return fail(std::string("inconsistent expert count: tensor ") + t->name + " has " +
|
|
std::to_string(ne2) + ", expected " + std::to_string(n_expert));
|
|
}
|
|
}
|
|
if (n_bound == 0) return fail("no MoE expert tensors captured — is this a MoE model?");
|
|
|
|
// Computed before init consumes `layers`: the dense split needs the captured expert
|
|
// tensor names and the gguf sizes together.
|
|
std::vector<uint64_t> dense_bytes;
|
|
if (route_trace) dense_bytes = dense_bytes_per_layer(offs, layers, im.n_layer);
|
|
|
|
// Anonymous dense-weights mode: hand the streamer the dense (non-expert) model weights to
|
|
// read into anon buffers. The list is every captured weight leaf that IS a gguf tensor
|
|
// (dropping graph inputs and KV, which share the leaf shape) and is NOT one of the streamed
|
|
// experts. Built before init consumes `layers`. Only this mode needs them; the others ignore
|
|
// an empty list.
|
|
if (cfg.moe.dense_weights == DenseWeightsMode::Anonymous || cfg.moe.dense_weights == DenseWeightsMode::Pinned) {
|
|
const std::unordered_set<std::string> expert_names = expert_tensor_names(layers);
|
|
std::vector<DenseTensorRef> dense;
|
|
for (const auto & kv : im.hook->captured_weights()) {
|
|
const std::string & name = kv.first;
|
|
if (expert_names.count(name)) continue;
|
|
auto off = offs.off_by_name.find(name);
|
|
auto sz = offs.size_by_name.find(name);
|
|
if (off == offs.off_by_name.end() || sz == offs.size_by_name.end()) continue; // not a file tensor
|
|
DenseTensorRef d;
|
|
d.tensor = kv.second;
|
|
d.file_off = off->second;
|
|
d.size = sz->second;
|
|
dense.push_back(d);
|
|
}
|
|
im.source.set_dense_tensors(std::move(dense));
|
|
}
|
|
|
|
if (!im.source.init(cfg.model_path, n_expert, std::move(layers), cfg.moe))
|
|
return fail("expert stream source init failed");
|
|
im.hook->set_source(&im.source);
|
|
|
|
if (route_trace) {
|
|
im.route_trace = route_trace;
|
|
im.hook->set_trace(true);
|
|
|
|
RouteTraceStatic st;
|
|
st.model = cfg.model_path;
|
|
st.arch = im.arch;
|
|
st.n_layer = im.n_layer;
|
|
st.n_expert = n_expert;
|
|
// The effective top-k: an override IS the applied width, otherwise the model's own.
|
|
st.n_expert_used = cfg.n_expert_used;
|
|
if (st.n_expert_used <= 0) {
|
|
const GgufModelInfo & info = gguf();
|
|
st.n_expert_used = info.ok ? info.n_expert_used : 0;
|
|
}
|
|
st.dense_bytes_per_layer = std::move(dense_bytes);
|
|
st.expert_bytes_per_layer.resize((size_t) im.n_layer);
|
|
for (int il = 0; il < im.n_layer; ++il)
|
|
st.expert_bytes_per_layer[(size_t) il] = im.source.expert_bytes(il);
|
|
route_trace->on_static(st);
|
|
}
|
|
|
|
if (io_trace) {
|
|
im.io_trace = io_trace;
|
|
im.source.set_io_trace(true);
|
|
}
|
|
|
|
if (cfg.moe.overlap) {
|
|
#ifdef BMOE_HAVE_EXPERT_READY_HOOK
|
|
im.source.enable_overlap_hook();
|
|
#else
|
|
return fail("--overlap requires the bmoe llama.cpp fork (expert-ready hook not compiled in)");
|
|
#endif
|
|
}
|
|
|
|
llama_memory_clear(llama_get_memory(ctx), true); // discard warm-up KV
|
|
}
|
|
|
|
// Decode traces. Outside the streaming block on purpose: the compute trace measures the graph,
|
|
// which exists with or without the streamer, so a dense mmap baseline can be traced and
|
|
// compared. The I/O trace was armed above (it needs the source) and only reports here.
|
|
if (compute_trace || im.io_trace) {
|
|
DecodeTraceStatic st;
|
|
st.model = cfg.model_path;
|
|
st.arch = im.arch;
|
|
st.n_layer = im.n_layer;
|
|
st.n_threads = cfg.n_threads;
|
|
st.io_threads = cfg.moe.enabled ? cfg.moe.io_threads : 0;
|
|
st.o_direct = cfg.moe.enabled && cfg.moe.o_direct;
|
|
st.overlap = cfg.moe.enabled && cfg.moe.overlap;
|
|
if (compute_trace) {
|
|
im.compute_trace = compute_trace;
|
|
im.hook->set_compute_trace(true, cfg.compute_trace_layers);
|
|
compute_trace->on_static(st);
|
|
}
|
|
if (im.io_trace) im.io_trace->on_static(st);
|
|
}
|
|
|
|
// What this session IS, for any metrics sink that will describe what it DOES. Built here, where
|
|
// every fact is resolved: cache_mb in particular is what the streamer settled on, which under
|
|
// auto-sizing is a number no flag ever mentioned.
|
|
{
|
|
RunInfo & ri = im.info;
|
|
const std::string & p = cfg.model_path;
|
|
const size_t slash = p.find_last_of("/\\");
|
|
ri.model = slash == std::string::npos ? p : p.substr(slash + 1);
|
|
ri.arch = im.arch;
|
|
ri.n_layer = im.n_layer;
|
|
ri.n_threads = cfg.n_threads;
|
|
ri.n_ctx = cfg.n_ctx;
|
|
ri.moe_stream = cfg.moe.enabled;
|
|
ri.cache_auto = cfg.moe.cache_auto;
|
|
ri.cache_ceil_mb = cfg.moe.cache_ceil_mb;
|
|
ri.force_cache = cfg.moe.force_cache;
|
|
ri.io_threads = cfg.moe.enabled ? cfg.moe.io_threads : 0;
|
|
ri.o_direct = cfg.moe.enabled && cfg.moe.o_direct;
|
|
ri.overlap = cfg.moe.enabled && cfg.moe.overlap;
|
|
ri.prefetch_layers = cfg.moe.enabled ? cfg.moe.prefetch_layers : 0;
|
|
ri.drop_cold_frac = cfg.moe.enabled ? cfg.moe.drop_cold_frac : 0.0f;
|
|
// The CSV keeps the two familiar flags, derived from the resolved dense-weights policy.
|
|
ri.dense_weights = cfg.moe.dense_weights == DenseWeightsMode::Mmap ? "mmap"
|
|
: cfg.moe.dense_weights == DenseWeightsMode::Anonymous ? "anon"
|
|
: cfg.moe.dense_weights == DenseWeightsMode::Pinned ? "ahwb"
|
|
: "warm";
|
|
if (cfg.moe.enabled) {
|
|
const IExpertSource::Stats st = im.source.stats();
|
|
ri.cache_mb = (int) (st.cache_budget_bytes / (1024ull * 1024ull));
|
|
}
|
|
// The EFFECTIVE top-k: an override IS the applied width, otherwise the model's own. Same
|
|
// resolution the route trace does, and worth a header read — a run whose top-k is unknown
|
|
// cannot be compared against one whose top-k differs, which is most of the point.
|
|
ri.n_expert_used = cfg.n_expert_used;
|
|
if (ri.n_expert_used <= 0 || ri.n_expert == 0) {
|
|
const GgufModelInfo & mi = gguf();
|
|
if (mi.ok) {
|
|
ri.n_expert = mi.n_expert;
|
|
if (ri.n_expert_used <= 0) ri.n_expert_used = mi.n_expert_used;
|
|
}
|
|
}
|
|
}
|
|
|
|
im.load_seconds = secs(t_load0, clock_t_::now());
|
|
return self;
|
|
}
|
|
|
|
RunResult Session::generate(const GenerateRequest & req,
|
|
const std::function<void(const TokenMetrics &)> & on_token,
|
|
IMetricsSink * sink) {
|
|
Impl & im = *impl_;
|
|
const MoeStreamConfig & moe = im.cfg.moe;
|
|
llama_context * ctx = im.ctx.get();
|
|
|
|
// Before anything is written about what this run did, say what it was.
|
|
if (sink && !im.info_sent) {
|
|
sink->on_run_info(im.info);
|
|
im.info_sent = true;
|
|
}
|
|
|
|
// Fresh cancel latch for this generation; a stale request from a prior aborted call must
|
|
// not carry over. (cancel() sets it; the abort callback reads it.)
|
|
im.cancel_requested.store(false, std::memory_order_relaxed);
|
|
|
|
RunResult res;
|
|
auto fail = [&](std::string msg) {
|
|
res.ok = false;
|
|
res.error = std::move(msg);
|
|
return res;
|
|
};
|
|
|
|
// clear_kv = "new chat": drop the KV and the engine-held conversation. Otherwise this turn
|
|
// continues the conversation, reusing the KV prefix already decoded from earlier turns.
|
|
if (req.clear_kv) {
|
|
llama_memory_clear(llama_get_memory(ctx), true);
|
|
im.chat_history.clear();
|
|
im.kv_tokens.clear();
|
|
// A new chat resets the sampler RNG, so a fixed seed reproduces the same transcript from a
|
|
// fresh conversation. A continued turn (clear_kv=false) keeps the stream going, matching the
|
|
// KV it decodes against.
|
|
if (im.smpl) llama_sampler_reset(im.smpl);
|
|
}
|
|
|
|
// Format the prompt. With chat on, render the model's OWN chat template (real Jinja) over the
|
|
// WHOLE conversation so far, and set up reasoning parsing so a thinking model's internal
|
|
// reasoning is stripped from the shown answer. req.think drives enable_thinking, per prompt.
|
|
std::string prompt = req.prompt;
|
|
bool chat_on = im.chat_on;
|
|
bool history_pushed = false; // did we append this turn's user message to chat_history?
|
|
bool prefilled_answer = false; // closed the reasoning span in the prompt, so skip reasoning parse
|
|
common_chat_parser_params parse_params;
|
|
if (chat_on) {
|
|
try {
|
|
common_chat_msg user_msg;
|
|
user_msg.role = "user";
|
|
user_msg.content = req.prompt;
|
|
im.chat_history.push_back(user_msg);
|
|
history_pushed = true;
|
|
|
|
common_chat_templates_inputs inputs;
|
|
inputs.messages = im.chat_history; // the full conversation, not just this turn
|
|
inputs.add_generation_prompt = true;
|
|
inputs.use_jinja = true;
|
|
inputs.enable_thinking = req.think;
|
|
// AUTO is what bakes reasoning-stripping into the generated parser grammar. It is set
|
|
// here, before apply — the field defaults to NONE, which produces a content-only
|
|
// grammar that leaves <think> markers in the answer no matter how the parse is wired.
|
|
inputs.reasoning_format = COMMON_REASONING_FORMAT_AUTO;
|
|
|
|
// Many templates never read enable_thinking (LFM2.5 among them): the flag reaches the
|
|
// jinja context, is discarded, and the model reasons anyway — the setting silently does
|
|
// nothing. For those, close the reasoning span in the prompt instead, so the model
|
|
// resumes at the first token of its answer with the reasoning already behind it.
|
|
//
|
|
// The span is rendered by llama.cpp's own handler for this template, so no marker for
|
|
// any family — harmony's primed final channel included — is spelled out here. Which
|
|
// models need this was measured at open(), not assumed. See thinking_control.h.
|
|
if (!req.think && im.think_ctl == ThinkControl::Prefill) {
|
|
detail::add_no_think_prefill(inputs);
|
|
prefilled_answer = true;
|
|
}
|
|
|
|
common_chat_params cp = common_chat_templates_apply(im.chat_tmpls.get(), inputs);
|
|
prompt = cp.prompt;
|
|
parse_params = detail::build_parse_params(cp);
|
|
} catch (const std::exception & e) {
|
|
std::fprintf(stderr, "bmoe: chat template apply failed (%s); using raw prompt\n", e.what());
|
|
if (history_pushed) {
|
|
im.chat_history.pop_back();
|
|
history_pushed = false;
|
|
}
|
|
chat_on = false;
|
|
}
|
|
}
|
|
|
|
std::vector<llama_token> tokens(prompt.size() + 8);
|
|
int n_prompt = llama_tokenize(im.vocab, prompt.c_str(), (int) prompt.size(), tokens.data(), (int) tokens.size(),
|
|
/*add_special*/ true, /*parse_special*/ true);
|
|
if (n_prompt < 0) {
|
|
tokens.resize(-n_prompt);
|
|
n_prompt = llama_tokenize(im.vocab, prompt.c_str(), (int) prompt.size(), tokens.data(), (int) tokens.size(),
|
|
true, true);
|
|
}
|
|
if (n_prompt < 1) return fail("empty prompt after tokenization");
|
|
tokens.resize(n_prompt);
|
|
if (n_prompt + req.n_predict + 8 > im.cfg.n_ctx)
|
|
return fail("prompt + n_predict exceeds the session n_ctx (" + std::to_string(im.cfg.n_ctx) +
|
|
"); open the session with a larger n_ctx");
|
|
|
|
// The text to surface: with chat on, parse the raw output so a reasoning model's internal
|
|
// thinking is separated from the answer. The answer is shown inline; the reasoning is handed to
|
|
// the UI as a distinct thinking block rather than dropped, so a Thinking-on run does not sit on a
|
|
// blank screen while the model reasons. Generation always uses the raw tokens.
|
|
struct ShownView {
|
|
std::string content; // the answer, reasoning stripped
|
|
std::string reasoning; // the thinking span, empty unless the parser split one out
|
|
};
|
|
auto shown_view = [&](const std::string & raw, bool partial) -> ShownView {
|
|
if (!chat_on) return {raw, ""};
|
|
// A prefilled turn resumes mid-answer: the reasoning span and the turn header the parser
|
|
// anchors on are already in the prompt, not in the stream, so there is nothing to strip —
|
|
// the raw stream already IS the answer. (If a model reasons anyway despite the closed span,
|
|
// that reasoning surfaces verbatim rather than being cut out here. Hiding it would only
|
|
// disguise a model this mechanism does not work on; the honest report is ThinkControl::None.)
|
|
if (prefilled_answer) return {raw, ""};
|
|
try {
|
|
common_chat_msg msg = common_chat_parse(raw, partial, parse_params);
|
|
return {msg.content, msg.reasoning_content};
|
|
} catch (const std::exception & e) {
|
|
detail::warn_parse_failed_once(e.what());
|
|
return {raw, ""};
|
|
}
|
|
};
|
|
|
|
// Reuse the KV prefix already decoded from earlier turns (chat mode only): find how many
|
|
// leading tokens still match the cache, drop the divergent tail, and prefill only the suffix.
|
|
// Keeping at least one token to decode means a turn is never a no-op. clear_kv leaves
|
|
// kv_tokens empty, so n_common = 0 and this reduces to a full prefill — the one-shot path the
|
|
// byte-identity gates exercise stays unchanged.
|
|
size_t n_common = 0;
|
|
if (chat_on && !im.kv_tokens.empty()) {
|
|
const size_t max_common = tokens.size() > 0 ? tokens.size() - 1 : 0;
|
|
while (n_common < im.kv_tokens.size() && n_common < max_common && im.kv_tokens[n_common] == tokens[n_common])
|
|
++n_common;
|
|
if (n_common < im.kv_tokens.size()) {
|
|
// SWA-style memory (e.g. Gemma) can refuse a partial removal; fall back to a full
|
|
// re-prefill in that case rather than continuing from an inconsistent cache.
|
|
if (!llama_memory_seq_rm(llama_get_memory(ctx), 0, (llama_pos) n_common, -1)) {
|
|
llama_memory_clear(llama_get_memory(ctx), true);
|
|
n_common = 0;
|
|
}
|
|
im.kv_tokens.resize(n_common);
|
|
}
|
|
}
|
|
|
|
// Roll this turn back to the state before it started: drop the KV added this turn, forget the
|
|
// tokens we fed, and un-append the user message. Used on cancel so prior turns stay usable.
|
|
auto rollback_turn = [&]() {
|
|
if (chat_on) {
|
|
if (!llama_memory_seq_rm(llama_get_memory(ctx), 0, (llama_pos) n_common, -1))
|
|
llama_memory_clear(llama_get_memory(ctx), true);
|
|
im.kv_tokens.resize(n_common);
|
|
if (history_pushed) {
|
|
im.chat_history.pop_back();
|
|
history_pushed = false;
|
|
}
|
|
} else {
|
|
llama_memory_clear(llama_get_memory(ctx), true);
|
|
}
|
|
};
|
|
|
|
// Route trace: frame one decode's rows, then hand them to the sink once it has returned —
|
|
// never from inside the callback, which runs on a compute thread mid-graph. `base_pos` is
|
|
// the context position of the batch's first token, so prefill rows carry real step numbers.
|
|
// The frame the I/O rows are stamped with at flush; the other traces carry their own.
|
|
int trace_phase = 0, trace_step = 0;
|
|
auto trace_begin = [&](int base_pos, int n_tokens, int phase) {
|
|
// Not a trace concern, but the same per-decode frame: the drop policy is decode-only
|
|
// unless armed for prefill, so it has to be told which phase this batch is.
|
|
im.hook->set_batch_phase(phase);
|
|
if (im.route_trace) im.hook->begin_trace_batch(base_pos, n_tokens, phase, im.turn);
|
|
// A node is computed once for the whole batch, not per token, so a prefill chunk's graph is
|
|
// attributed to its last position rather than pretending to split across the chunk.
|
|
if (im.compute_trace) im.hook->begin_compute_batch(base_pos + n_tokens - 1, phase, im.turn);
|
|
trace_phase = phase;
|
|
trace_step = base_pos + n_tokens - 1;
|
|
};
|
|
auto trace_flush = [&]() {
|
|
// Close the compute trace's dangling interval FIRST: at layer granularity the "post" row
|
|
// is charged the wall since the last boundary, and everything trace_flush does before the
|
|
// close would be billed to the LM head.
|
|
if (im.compute_trace) im.hook->end_compute_batch();
|
|
if (im.route_trace) {
|
|
im.hook->end_trace_batch();
|
|
std::vector<RouteTraceRow> & rows = im.hook->trace_rows();
|
|
if (!rows.empty()) im.route_trace->on_rows(rows.data(), rows.size());
|
|
rows.clear();
|
|
}
|
|
if (im.compute_trace) {
|
|
std::vector<ComputeTraceRow> & rows = im.hook->compute_rows();
|
|
if (!rows.empty()) im.compute_trace->on_rows(rows.data(), rows.size());
|
|
rows.clear();
|
|
}
|
|
if (im.io_trace) {
|
|
// The reads carry no frame of their own — a lane does not know which token it serves —
|
|
// so stamp them with the decode they were drained after.
|
|
im.source.take_io_trace_rows(im.io_rows_scratch);
|
|
for (IoTraceRow & r : im.io_rows_scratch) {
|
|
r.turn = im.turn;
|
|
r.phase = trace_phase;
|
|
r.step = trace_step;
|
|
}
|
|
if (!im.io_rows_scratch.empty()) im.io_trace->on_rows(im.io_rows_scratch.data(), im.io_rows_scratch.size());
|
|
im.io_rows_scratch.clear();
|
|
}
|
|
};
|
|
|
|
// ── prefill (chunked by n_batch; positions auto-continue from the reused prefix) ──
|
|
const auto t_prefill0 = clock_t_::now();
|
|
for (int i = (int) n_common; i < n_prompt; i += im.cfg.n_batch) {
|
|
const int chunk = std::min(im.cfg.n_batch, n_prompt - i);
|
|
llama_batch pf = llama_batch_get_one(tokens.data() + i, chunk);
|
|
trace_begin(i, chunk, /*phase*/ 0);
|
|
if (llama_decode(ctx, pf) != 0) {
|
|
if (im.cancel_requested.load(std::memory_order_relaxed)) {
|
|
rollback_turn();
|
|
res.ok = true;
|
|
res.cancelled = true;
|
|
return res;
|
|
}
|
|
if (moe.overlap && im.source.fatal()) return fail("expert stream I/O failed during overlap prefill");
|
|
return fail("prefill decode failed");
|
|
}
|
|
trace_flush();
|
|
}
|
|
// The suffix is now in the KV; record it so the next turn can diff against it.
|
|
if (chat_on)
|
|
for (int i = (int) n_common; i < n_prompt; ++i)
|
|
im.kv_tokens.push_back(tokens[i]);
|
|
const double prefill_seconds = secs(t_prefill0, clock_t_::now());
|
|
const float * logits = llama_get_logits_ith(ctx, -1);
|
|
|
|
// ── greedy generation ──
|
|
res.ok = true;
|
|
std::string gen;
|
|
int n_gen = 0;
|
|
double gen_seconds = 0.0;
|
|
|
|
// Baseline snapshot taken AFTER prefill: the summary reports the generation phase only,
|
|
// from real per-token deltas (prefill routes near the whole bank, so folding it into a
|
|
// per-token average would badly inflate the flash-I/O figure). In a warm session these
|
|
// counters carry the prior prompts' totals; the deltas make each prompt self-relative.
|
|
GenTally tally;
|
|
tally.overlap = moe.overlap;
|
|
tally.n_threads = im.cfg.n_threads;
|
|
if (moe.enabled) {
|
|
const IExpertSource::Stats st0 = im.source.stats();
|
|
tally.prev_bytes = (long long) st0.read_bytes;
|
|
tally.prev_io_s = st0.read_seconds;
|
|
tally.prev_mgmt_s = st0.mgmt_seconds;
|
|
tally.prev_stall_s = st0.stall_seconds;
|
|
}
|
|
long long prev_spec_bytes = moe.enabled ? (long long) im.source.stats().spec_read_bytes : 0;
|
|
long long prev_spec_experts = moe.enabled ? im.source.stats().spec_experts : 0;
|
|
long long prev_spec_useful = moe.enabled ? im.source.stats().spec_useful : 0;
|
|
// Taken after prefill, so the drop counters describe generation — the phase the policy is armed
|
|
// for and the one the tok/s number is about.
|
|
const long long prev_routed = im.hook->experts_routed();
|
|
const long long prev_dropped = im.hook->experts_dropped();
|
|
|
|
for (int t = 0; t < req.n_predict; ++t) {
|
|
// Greedy stays argmax (byte-identical to the resident reference the gates check); with a
|
|
// sampling chain, draw from the context's last-position logits, which llama_sampler_sample
|
|
// reads at index -1 — the same logits argmax would have read.
|
|
llama_token tok = im.smpl ? llama_sampler_sample(im.smpl, ctx, -1) : argmax(logits, im.n_vocab);
|
|
if (llama_vocab_is_eog(im.vocab, tok)) break;
|
|
|
|
char piece[256];
|
|
int np = llama_token_to_piece(im.vocab, tok, piece, sizeof(piece), 0, true);
|
|
std::string delta = np > 0 ? std::string(piece, np) : std::string();
|
|
gen += delta;
|
|
|
|
// Bracket ONLY the decode: major faults and CPU-time deltas here decompose this token's
|
|
// compute residual into flash-fault stalls vs. genuine (or throttled) computation.
|
|
const uint64_t f0 = pio::major_faults();
|
|
const double c0 = pio::process_cpu_seconds();
|
|
auto s0 = clock_t_::now();
|
|
llama_batch step = llama_batch_get_one(&tok, 1);
|
|
trace_begin(n_prompt + t, /*n_tokens*/ 1, /*phase*/ 1); // decodes at the position after the prompt
|
|
int dec = llama_decode(ctx, step);
|
|
auto s1 = clock_t_::now();
|
|
const uint64_t f1 = pio::major_faults();
|
|
const double c1 = pio::process_cpu_seconds();
|
|
if (dec != 0) {
|
|
if (im.cancel_requested.load(std::memory_order_relaxed)) {
|
|
res.cancelled = true;
|
|
break;
|
|
}
|
|
if (moe.overlap && im.source.fatal()) return fail("expert stream I/O failed during overlap decode");
|
|
return fail("decode failed during generation");
|
|
}
|
|
trace_flush(); // outside the s0..s1 bracket: the trace's own writes must not bill wall_ms
|
|
logits = llama_get_logits_ith(ctx, -1);
|
|
if (chat_on) im.kv_tokens.push_back(tok); // this token is now in the KV
|
|
|
|
++n_gen;
|
|
double wall = secs(s0, s1);
|
|
gen_seconds += wall;
|
|
|
|
TokenMetrics m;
|
|
m.step = n_gen;
|
|
m.steps = req.n_predict;
|
|
m.piece = delta;
|
|
{
|
|
ShownView sv = shown_view(gen, /*partial*/ true);
|
|
m.text = std::move(sv.content);
|
|
m.reasoning = std::move(sv.reasoning);
|
|
}
|
|
const IExpertSource::Stats st = moe.enabled ? im.source.stats() : IExpertSource::Stats{};
|
|
tally.record(m, wall, f1 - f0, c1 - c0, im.turn, moe.enabled ? &st : nullptr);
|
|
if (on_token) on_token(m);
|
|
if (sink) sink->on_token(m);
|
|
}
|
|
|
|
++im.turn; // this turn is written; label the next one apart
|
|
|
|
// ── summary ──
|
|
RunSummary & s = res.summary;
|
|
s.n_generated = n_gen;
|
|
s.gen_seconds = gen_seconds;
|
|
s.s_per_token = n_gen ? gen_seconds / n_gen : 0.0;
|
|
s.tokens_per_second = gen_seconds > 0 ? n_gen / gen_seconds : 0.0;
|
|
s.n_prompt = n_prompt - (int) n_common; // tokens actually prefilled this turn (after KV reuse)
|
|
s.n_past = chat_on ? (int) im.kv_tokens.size() : n_prompt + n_gen; // total context length now
|
|
s.load_seconds = im.load_seconds;
|
|
s.prefill_seconds = prefill_seconds;
|
|
s.majflt_per_token = n_gen ? (double) tally.majflt / n_gen : 0.0;
|
|
s.cpu_s_per_token = n_gen ? tally.cpu_seconds / n_gen : 0.0;
|
|
if (moe.enabled) {
|
|
IExpertSource::Stats st = im.source.stats();
|
|
s.moe_read_mib = tally.read_bytes / (1024.0 * 1024.0);
|
|
s.moe_io_seconds = tally.io_seconds;
|
|
s.moe_io_s_per_token = n_gen ? tally.io_seconds / n_gen : 0.0;
|
|
s.moe_mgmt_s_per_token = n_gen ? tally.mgmt_seconds / n_gen : 0.0;
|
|
s.moe_stall_s_per_token = n_gen ? tally.stall_seconds / n_gen : 0.0;
|
|
s.moe_compute_s_per_token =
|
|
s.s_per_token - (moe.overlap ? s.moe_stall_s_per_token : s.moe_io_s_per_token) - s.moe_mgmt_s_per_token;
|
|
if (s.moe_compute_s_per_token < 0) s.moe_compute_s_per_token = 0;
|
|
s.cache_hit_pct = st.cache_lookups > 0 ? 100.0 * st.cache_hits / st.cache_lookups : -1.0;
|
|
s.cache_resident_mib = st.cache_resident_bytes / (1024.0 * 1024.0);
|
|
s.cache_budget_mib = st.cache_budget_bytes / (1024.0 * 1024.0);
|
|
s.cache_resizes = st.cache_resizes;
|
|
s.token_demand_mib = st.token_demand_bytes / (1024.0 * 1024.0);
|
|
s.layer_demand_mib = st.layer_demand_bytes / (1024.0 * 1024.0);
|
|
s.moe_spec_read_mib = ((long long) st.spec_read_bytes - prev_spec_bytes) / (1024.0 * 1024.0);
|
|
s.moe_spec_experts = st.spec_experts - prev_spec_experts;
|
|
s.moe_spec_useful = st.spec_useful - prev_spec_useful;
|
|
}
|
|
s.experts_routed = im.hook->experts_routed() - prev_routed;
|
|
s.experts_dropped = im.hook->experts_dropped() - prev_dropped;
|
|
if (sink) sink->on_summary(s);
|
|
|
|
{
|
|
ShownView sv = shown_view(gen, /*partial*/ false);
|
|
res.generated_text = std::move(sv.content);
|
|
res.reasoning_text = std::move(sv.reasoning);
|
|
}
|
|
|
|
if (res.cancelled) {
|
|
// Undo the whole turn (KV, fed tokens, and the pushed user message) so the conversation
|
|
// is left exactly as it was before this prompt and stays continuable.
|
|
rollback_turn();
|
|
} else if (chat_on) {
|
|
// Commit the assistant turn to the running conversation. Parsing separates a thinking
|
|
// model's reasoning from the answer; the next turn re-renders history from these messages.
|
|
common_chat_msg assistant;
|
|
if (prefilled_answer) {
|
|
// Prefilled turn: no turn header in the stream to parse; the generation is the answer
|
|
// verbatim. It is committed without the prefill, so history holds a normal assistant
|
|
// message and the next turn re-renders cleanly whatever this turn's think setting was.
|
|
assistant.content = gen;
|
|
} else {
|
|
try {
|
|
assistant = common_chat_parse(gen, /*is_partial*/ false, parse_params);
|
|
} catch (const std::exception & e) {
|
|
detail::warn_parse_failed_once(e.what());
|
|
assistant.content = gen;
|
|
}
|
|
}
|
|
assistant.role = "assistant";
|
|
im.chat_history.push_back(assistant);
|
|
}
|
|
return res;
|
|
}
|
|
|
|
} // namespace bmoe
|