diff --git a/CHANGELOG.md b/CHANGELOG.md index f4189a0..652b781 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -166,6 +166,76 @@ Semantic Versioning. width (5.59 at draft 2 and 4.38 at draft 3 against 5.82–6.14 baseline) because the shipping configuration is compute-bound — `stall_s/tok` is 0.025–0.027 whether speculation is on or off — while the read set widens 35–54% and CPU per token rises 28–67%. The flag stays off. +- **`--route-ahead N` — commit decode routing to the N-layers-early prediction (experimental, + lossy, off by default).** Every prefetch lives under the same ceiling: layer L's routing needs + layer L−1's output, so any predictor working earlier is approximate and every speculated read + can miss. This flag inverts the bet — the expert selection of decode layer L is *replaced* by + the ranking layer L's own gate produced on the hidden state N layers back in the same forward + pass, so the selection is known N full layers of compute early and a prefetch of it could never + be wrong (the training-free cousin of Pre-gated MoE, ISCA'24). The router still computes: its + logits give the substituted experts their true renormalized weights via the graph's own weight + chain, and each layer's gate input feeds the prediction for layer L+N. Layers 0..N−1, prefill, + the first decode token and anything unreadable route normally and are counted as passed + through. The `moe-route-ahead:` report line and the `route_ahead=` CSV preamble key record how + many routings were committed and how far the committed selection sat from the router's own + choice — the measured perturbation the run generated under. Mutually exclusive with + `--predict-log` and both prefetchers; composes with the cache and expert dropping. With the + LRU cache on, the committed selection is also acted on: layer L+N's ids are handed to the + speculative read path the moment they are fixed (at layer L's load), uncapped — unlike every + predictor before them these reads can never be wasted, so the `moe-prefetch: [route-ahead]` + line reports a useful fraction that sits at ~100% by construction — and the issue list is + drop-aware: committed experts predicted below the `--drop-cold-experts` threshold are left + cold on purpose so the commit-time drop discards them unread, exactly as it does the router's + own tail (measured on device before this filter: uncapped early reads un-dropped everything, + 3x the flash per token and half the tok/s at depth 4). The prediction itself runs + on the barrier-less path predictive prefetch built — pointers stashed in the ask pass, the row + copied at the topk callback, the GEMV on the shared worker thread, the sampled fresh-gate + watchdog validating the read — so no graph barrier is added and the eval thread pays only a + row copy per layer; a watchdog trip disarms the policy out loud, and a layer whose ranking is + not ready when its topk fires keeps the router's own choice. On the streamer side, committed + speculation is exempt from the destructive quiesce built for guessing predictors: a loading + layer adopts its own committed jobs (finishing them instead of discarding and re-reading), and + other layers' committed reads survive in the queue — without this, every early read at depth 2 + was destroyed before use. Host A/B under `--overlap` (fixed cache, same prompt): N=2 cuts the + overlap stall 0.060 → 0.010 s/token (97% of early reads useful) for +20% tok/s end to end, + with the generation still correct. **Device: components proven, end-to-end owed.** An I/O trace of every physical read + (decode-only, equal length) shows the committed routing reads what the baseline reads — 133.6 + vs 134.4 MiB/token, the same 4% redundancy, zero speculative-then-demand pairs — so nothing is + fetched twice, and the overlap stall falls consistently from 0.038-0.059 to 0.006-0.007 + s/token. **Device throughput, measured: +21%.** Interleaved cells (baseline / N=2, two passes, + same prompt and length, so thermal drift hits both equally) put N=2 at 5.29 tok/s against 4.36 + for the baseline, with the overlap stall down 0.067 → 0.010 s/token and the compute residual + essentially unchanged (0.135 → 0.146) — the win is the stall, not a trade. The same build + measures only +4% inside the demo app, and the reason is not the engine: with the app in + foreground the device caps its six main cores at 1.9 GHz (2188800 → 1900800 Hz, measured), + which slows the compute the flag is trying to keep busy and shrinks the share of the token + that was stall. Two real costs were found by instrumenting rather than arguing, and both are fixed: + the issue path took the I/O lanes' mutex once per expert (~120 acquisitions a token; the same + run with one lane did identical work at a quarter of the compute time), now 0.9 ms/token; and + the gate GEMV was bandwidth-bound rather than compute-bound — ranking 256 experts re-reads a + ~4 MB gate matrix per layer — now served from an int8 mirror with a quantized activation + (aarch64 only; on x86 the float path is already vectorized and the detour costs more than it + saves), 19 → 6 ms/token with the committed selection agreeing on 74.5% of slots against 74.7% + for the exact weights. What is still missing is a clean throughput pair: after a long + benchmarking session the phone's own baseline swung between 0.79 and 4.54 tok/s in the same + cell — which is why the figure above comes from interleaved cells rather than absolutes. Every + eval-thread cost the flag adds now reports itself in the CSV and the summary (`ra_issue_ms`, + `ra_wd_ms`, `adopt_ms` inside mgmt, `drain_ms` inside compute, plus cache `evictions`/`rereads`), + because five successive explanations of this feature's device cost were deduced from a residual + and all five were wrong. One structural cost remains and is documented as the next step: the graph + computes its own router matmul regardless, so this implementation runs two gate matmuls per + layer where the design wants one substituted — removing the second one means changing the + graph, i.e. the fork that already exists for the expert-ready hook. Quality is measured, not + assumed: the committed generations match the baseline's on a four-prompt objective battery + (4/4 correct in both, one answer byte-identical), on a 512-token essay, and on a second model + of a different generation and quantization (+22% there); output is deterministic across + repeated runs, which expert dropping is not. It is refused alongside self-speculation, and that + exclusion is measured rather than reasoned: a verify decode is several positions wide, so the + policy declines to commit on every one of them while still running its prediction and issuing + early reads that have become ordinary speculation (a combined run committed 0 routings, passed + 249 through, and dropped its speculative usefulness from 100% to 81%). Cost without commitment + is worse than either feature alone, so `validate()` rejects the pair instead of silently + charging for it. See `docs/route-ahead.md`. ## [0.18.1] - 2026-07-29 diff --git a/cli/main.cpp b/cli/main.cpp index 8de23e2..85ccaf5 100644 --- a/cli/main.cpp +++ b/cli/main.cpp @@ -436,6 +436,13 @@ static void print_usage(const char * argv0) { " it changes the output, and not reproducibly. Off by default.\n" " --drop-no-renorm do not rescale the surviving weights after a drop (A/B)\n" " --drop-in-prefill drop during prefill too (off: the cold cache makes it expensive)\n" + " --route-ahead N EXPERIMENTAL, LOSSY: commit decode routing to the prediction made\n" + " N layers earlier in the same forward pass (each layer's own gate\n" + " run on the hidden state N layers back). The router still computes\n" + " and gives the substituted experts their true renormalized weights;\n" + " a prefetch of a committed layer can then never miss. Changes the\n" + " output — this flag exists to measure that quality trade [0..8].\n" + " Excludes --predict-log / --predict-prefetch / --prefetch.\n" " --predict-log diagnostics: measure how much of each layer's routing could be\n" " known a layer early (the next layer's gate run on this layer's\n" " input), scored against the previous-token bet --prefetch makes.\n" @@ -647,6 +654,8 @@ int main(int argc, char ** argv) { cfg.moe.drop_renorm = false; else if (a == "--drop-in-prefill") cfg.moe.drop_prefill = true; + else if (a == "--route-ahead") + cfg.moe.route_ahead = std::atoi(next("--route-ahead")); else if (a == "--predict-log") cfg.moe.predict_log = true; else if (a == "--predict-prefetch") @@ -829,15 +838,31 @@ int main(int argc, char ** argv) { s.cache_resident_mib, s.cache_budget_mib); else std::printf("moe-cache: %.1f%% hit, resident %.1f MiB\n", s.cache_hit_pct, s.cache_resident_mib); + // Churn: a read of an entry the cache already held once. The bytes a routing needs are + // fixed, so this is where any surplus goes — and the number to compare across an A/B + // whose byte count moved. + if (s.cache_evictions > 0 || s.cache_rereads > 0) + std::printf("moe-cache: %lld evictions, %lld re-reads (%.1f/token) — bytes the cache had " + "already paid for once\n", + s.cache_evictions, s.cache_rereads, + s.n_generated ? (double) s.cache_rereads / s.n_generated : 0.0); } if (cfg.moe.overlap) std::printf("moe-overlap: stall %.3f s/token (flash reads overlapped with FFN compute)\n", s.moe_stall_s_per_token); - if (cfg.moe.prefetch_layers > 0 || cfg.moe.predict_prefetch) + // The named eval-thread waits, printed only when they cost something: the previous-batch + // drain lives inside the compute residual, the adoption wait inside mgmt. Either being + // large is a finding, not a footnote — both were invisible before they had meters. + if (s.moe_drain_s_per_token >= 0.0005 || s.moe_adopt_s_per_token >= 0.0005) + std::printf("moe-waits: drain %.3f s/token (inside compute), adopt %.3f s/token (inside mgmt)\n", + s.moe_drain_s_per_token, s.moe_adopt_s_per_token); + if (cfg.moe.prefetch_layers > 0 || cfg.moe.predict_prefetch || cfg.moe.route_ahead > 0) std::printf("moe-prefetch: %.1f MiB speculative, %lld/%lld experts useful (%.0f%%)%s\n", s.moe_spec_read_mib, s.moe_spec_useful, s.moe_spec_experts, s.moe_spec_experts > 0 ? 100.0 * s.moe_spec_useful / s.moe_spec_experts : 0.0, - cfg.moe.predict_prefetch ? " [stale-gate]" : ""); + cfg.moe.route_ahead > 0 ? " [route-ahead]" + : cfg.moe.predict_prefetch ? " [stale-gate]" + : ""); // How hard the policy actually bit. The flag sets a threshold, not a drop rate: what gets // discarded depends on what the cache held, so this is the only honest report of the trade // a given run made. @@ -846,6 +871,23 @@ int main(int argc, char ** argv) { s.experts_dropped, s.experts_routed, s.experts_routed > 0 ? 100.0 * s.experts_dropped / s.experts_routed : 0.0, (double) cfg.moe.drop_cold_frac); + // The agreement is the honest label for what the run just generated under: 100% minus it + // is the fraction of routed slots that went to an expert the router did not choose. + if (cfg.moe.route_ahead > 0) { + std::printf("moe-route-ahead: %d layers early — %lld routings committed, %lld passed through; " + "committed selection agreed with the router on %.1f%% of slots\n", + cfg.moe.route_ahead, s.route_ahead_overridden, s.route_ahead_passthrough, + s.route_ahead_slots > 0 ? 100.0 * s.route_ahead_hits / s.route_ahead_slots : 0.0); + // The prediction's own CPU, per token. On a host with spare cores this hides from the + // wall clock; on a phone it competes with the decode, so it is printed either way. + if (s.route_ahead_gemv_jobs > 0) + std::printf("moe-route-ahead: prediction %.1f ms/token of worker CPU (%lld GEMVs, %.3f ms each)" + " + %.1f ms/token issuing the reads + %.1f ms/token watchdog, both on the eval thread\n", + s.n_generated ? s.route_ahead_gemv_ns / 1e6 / s.n_generated : 0.0, s.route_ahead_gemv_jobs, + s.route_ahead_gemv_ns / 1e6 / s.route_ahead_gemv_jobs, + s.n_generated ? s.route_ahead_issue_ns / 1e6 / s.n_generated : 0.0, + s.n_generated ? s.route_ahead_wd_ns / 1e6 / s.n_generated : 0.0); + } if (cfg.moe.predict_log) print_predict_report(s); } return 0; diff --git a/core/include/bmoe/config.h b/core/include/bmoe/config.h index e4de597..b4a337e 100644 --- a/core/include/bmoe/config.h +++ b/core/include/bmoe/config.h @@ -172,6 +172,34 @@ struct MoeStreamConfig { // speculation, if any, is small (the stall a prefetch can remove is head-of-line only). int predict_spec_max = 2; + // ── route-ahead: commit to the prediction (lossy; opt-in; experimental) ─────────── + // With N > 0, decode routing at MoE layer L is REPLACED by the ranking layer L's own gate + // matrix produced on the hidden state N layers earlier in the same forward pass — the + // stale-gate prediction predict_log measures, acted on as the routing itself instead of as a + // prefetch hint. The router still computes: its logits give the substituted experts their + // true renormalized weights, and each layer's gate input feeds the prediction for layer L+N. + // Layers 0..N-1, prefill, the first decode token and any unreadable layer route normally. + // + // Why anyone would do this: a prediction that IS the routing cannot miss, and it is known a + // full N layers of compute early — the perfect-prefetch regime no honest predictor reaches + // (measured here: one layer of staleness costs ~3.5-4pp of slot agreement, and the union of + // four stale horizons still covers only 81-89% of misses). With the LRU cache on, the engine + // acts on that: the committed selection of layer L+N is handed to the speculative read path + // the moment it is fixed (at layer L's load), uncapped — those reads can never be wasted — + // so by the time layer L+N runs its experts are resident or already in flight. Without the + // cache the selection still commits but nothing is read early. The price is a quality + // perturbation: ~15-20% of slots route to a different expert than the router chose at N=1, + // and the error compounds through the residual stream. LOSSY by construction; changes the + // output. See docs/route-ahead.md. + // + // Mutually exclusive with the probe and both prefetchers: the probe would score a predictor + // against a routing this policy rewrote from that same predictor (a tautology), and a + // speculative prefetch would bet lanes on a future this policy has already fixed. Also + // mutually exclusive with self-speculation, for a measured reason rather than a conceptual + // one: a verify decode is several positions wide and this policy declines to commit on every + // one of them, so the pair pays the prediction and the early reads and commits nothing. + int route_ahead = 0; + // Test/debug only: complete each prefetch's speculative reads synchronously, on the eval // thread, before returning. This defeats the latency-hiding purpose (the reads no longer // overlap compute) but makes speculative integration deterministic, so the byte-identity @@ -189,6 +217,7 @@ struct MoeStreamConfig { static constexpr int cache_min_mb = 1500; // smallest non-pathological cache (see above) static constexpr int io_threads_max = 8; static constexpr int prefetch_layers_max = 8; + static constexpr int route_ahead_max = 8; // beyond this the staleness has no measured meaning }; // Where the draft tokens of a self-speculative step come from. The verify half of the loop is diff --git a/core/include/bmoe/expert_source.h b/core/include/bmoe/expert_source.h index eb8265d..8167f3d 100644 --- a/core/include/bmoe/expert_source.h +++ b/core/include/bmoe/expert_source.h @@ -78,6 +78,21 @@ public: long long spec_useful = 0; // prefetched experts that a later lookup actually hit uint64_t cache_budget_bytes = 0; // cache budget in force; fixed for the run once init sizes it long long cache_resizes = 0; // explicit set_cache_budget_mb() calls that moved the budget + // Cache churn. `evictions` is how many entries the budget forced out; `rereads` how many + // reads went to an entry that had been resident before — the cache paying for the same + // bytes twice. A prefetch cannot reduce what a routing needs (the ideal is the same + // bytes, earlier), so rereads is the only way a prefetch whose every read is USEFUL can + // still raise the byte count, and therefore the number to look at when it does. + long long evictions = 0; + long long rereads = 0; + // Eval-thread waits that are NOT part of the windows above, split out because each lives + // inside a different residual and hid there. `drain_wait` is the top of the async load: + // waiting for the previous layer's batch to finish before its jobs_/flags are reused — + // billed to neither io nor mgmt, so it sits in the compute residual. `adopt_wait` is the + // route-ahead adoption: the load waiting for its own committed speculative reads to + // complete before staging — inside the mgmt window, so this names its share of mgmt. + double drain_wait_seconds = 0.0; + double adopt_wait_seconds = 0.0; // ── residency telemetry (diagnostic) ── // Sampled fraction of the DENSE weights still in RAM, or -1 when not measured yet. Under the diff --git a/core/include/bmoe/metrics.h b/core/include/bmoe/metrics.h index cb7a943..4618d11 100644 --- a/core/include/bmoe/metrics.h +++ b/core/include/bmoe/metrics.h @@ -22,6 +22,16 @@ struct TokenMetrics { double mgmt_ms = 0.0; // cache-management time (vm commit + evict + LRU bookkeeping) this token double compute_ms = 0.0; // residual: serial wall - io - mgmt; overlap wall - stall - mgmt double stall_ms = 0.0; // overlap only: wall time the FFN kernel blocked on flash (0 when serial) + // Named eval-thread costs that otherwise hide inside a residual (all 0 when the feature that + // pays them is off). drain_ms sits inside compute_ms: the async load waiting out the previous + // layer's batch before reusing its flags. adopt_ms sits inside mgmt_ms: the route-ahead load + // waiting for its own committed speculative reads. ra_issue_ms and ra_wd_ms sit inside + // compute_ms: issuing the committed early reads, and the sampled fresh-gate watchdog. They + // exist because every wrong theory about this feature was a cost deduced from a residual. + double drain_ms = 0.0; + double adopt_ms = 0.0; + double ra_issue_ms = 0.0; + double ra_wd_ms = 0.0; // Wall time between the PREVIOUS token's decode and this one's: sampling, detokenization, // rendering the answer for a UI, this struct, and the sinks. None of it is inside wall_ms, so // none of it reaches the reported tok/s — which is exactly why it is worth a column. A change @@ -125,6 +135,15 @@ struct RunSummary { double cache_resident_mib = 0.0; double cache_budget_mib = 0.0; // the fixed cache budget the run used (explicit, or auto-sized at load) long long cache_resizes = 0; // runtime budget changes — now only an app's explicit set_cache_budget + // Cache churn (see IExpertSource::Stats): entries the budget forced out, and reads that went + // to an entry the cache had held before. Any prefetch that raises the byte count while + // claiming its reads are useful is doing it here. + long long cache_evictions = 0; + long long cache_rereads = 0; + // The named eval-thread waits (see TokenMetrics): the async load's previous-batch drain (part + // of the compute residual) and the route-ahead adoption wait (part of mgmt), per token. + double moe_drain_s_per_token = 0.0; + double moe_adopt_s_per_token = 0.0; // What one token actually demands of the cache, measured: the distinct expert bytes routed per // token. A cache below this can hold nothing between tokens; a cache far above it is buying @@ -195,6 +214,25 @@ struct RunSummary { PredictorStats predict_stale, predict_stale2, predict_prev, predict_self; std::vector predict_stale_by_layer, predict_prev_by_layer, predict_self_by_layer; long long predict_unscored = 0; // routings the stale-gate probe could not rank (see RouterHook) + + // Route-ahead (all zero unless MoeStreamConfig::route_ahead > 0): how many decode routings had + // their selection replaced by the N-layers-early prediction, how many eligible ones had no + // prediction to commit and kept the router's choice, and — over the overridden ones — how many + // slots the committed selection shared with what the router would have picked. hits/slots is + // the routing perturbation the run actually generated under; like the probe's counters these + // are session totals, not per-generation deltas. + long long route_ahead_overridden = 0; + long long route_ahead_passthrough = 0; + long long route_ahead_slots = 0; + long long route_ahead_hits = 0; + // The prediction's own CPU bill: worker nanoseconds spent on the gate GEMV, and how many + // rankings it produced. Reported per token because it is the one cost that hides from wall + // time wherever a spare core exists — and stops hiding on a phone, where the same core and + // the same DRAM serve the decode this is meant to accelerate. + long long route_ahead_gemv_ns = 0; + long long route_ahead_gemv_jobs = 0; + long long route_ahead_issue_ns = 0; // eval-thread cost of issuing the early reads + long long route_ahead_wd_ns = 0; // eval-thread cost of the sampled fresh-gate watchdog }; // What this run IS: the model and the configuration every row below it was produced under. @@ -232,6 +270,7 @@ struct RunInfo { bool overlap = false; bool io_two_wave = false; // two-wave batch publish (#118): first projection published early int prefetch_layers = 0; + int route_ahead = 0; // decode routing committed to the N-layers-early prediction (LOSSY) bool predict_prefetch = false; // stale-gate predictive prefetch (see MoeStreamConfig) bool predict_log = false; // the accuracy probe: costs a barrier and two GEMVs per layer int predict_spec_max = 0; // how many predicted misses a layer may speculate (0 = retention only) diff --git a/core/src/config.cpp b/core/src/config.cpp index dbe68c0..0eb44d3 100644 --- a/core/src/config.cpp +++ b/core/src/config.cpp @@ -118,6 +118,10 @@ ValidationResult validate(const RunConfig & cfg) { if (cfg.moe.predict_prefetch && !cfg.moe.enabled) { return fail("moe.predict_prefetch requires moe.enabled"); } + if (cfg.moe.route_ahead > 0 && !cfg.moe.enabled) { + return fail("moe.route_ahead requires moe.enabled: it attaches to the routing nodes the " + "streamer isolates, and off the streaming path there is nothing to commit to"); + } if (cfg.moe.enabled) { const MoeStreamConfig & m = cfg.moe; @@ -177,6 +181,29 @@ ValidationResult validate(const RunConfig & cfg) { "are two predictors for the same speculative read lanes, and running both " "doubles the speculated bytes for the same future."); } + if (m.route_ahead < 0 || m.route_ahead > MoeStreamConfig::route_ahead_max) { + return fail("moe.route_ahead must be in [0, " + std::to_string(MoeStreamConfig::route_ahead_max) + + "] (0 = off)"); + } + if (m.route_ahead > 0 && (m.predict_log || m.predict_prefetch || m.prefetch_layers > 0)) { + return fail("moe.route_ahead excludes predict_log, predict_prefetch and prefetch_layers: " + "the probe would grade a predictor against a routing rewritten from that same " + "predictor, and a speculative prefetch would bet lanes on a future route-ahead " + "has already fixed."); + } + // Measured, not theorised: with a draft source on, a verify decode is several positions wide + // and route-ahead declines every one of them — a run at draft 3 committed 0 routings and + // passed 249 through. It still pays for itself, though: the prediction GEMVs run, and its + // early reads become ordinary speculation that can now miss (81% useful against 100% when it + // commits). Cost with no commitment is worse than either feature alone, so the pair is + // refused rather than silently charged for. Making it work means committing the whole verify + // batch to one selection, which costs draft acceptance; see docs/route-ahead.md. + if (m.route_ahead > 0 && cfg.spec.enabled()) { + return fail("moe.route_ahead and self-speculative decoding are mutually exclusive: a " + "verify decode is several positions wide, so route-ahead declines to commit " + "on every one of them while still paying for its prediction and its early " + "reads. Turn off --mtp/--ngram, or turn off --route-ahead."); + } if (m.drop_cold_frac > 0.0f && !cache_on) { return fail("moe.drop_cold_frac requires the LRU cache (cache_mb > 0 or cache_auto): with the " "cache off every expert is a miss, so the policy stops being cache-aware and " diff --git a/core/src/engine/session.cpp b/core/src/engine/session.cpp index 4818a65..2136ea6 100644 --- a/core/src/engine/session.cpp +++ b/core/src/engine/session.cpp @@ -109,6 +109,10 @@ struct GenTally { double io_seconds = 0.0; double mgmt_seconds = 0.0; double stall_seconds = 0.0; + double drain_seconds = 0.0; + double adopt_seconds = 0.0; + double prev_drain_s = 0.0; + double prev_adopt_s = 0.0; uint64_t majflt = 0; double cpu_seconds = 0.0; @@ -166,15 +170,21 @@ struct GenTally { } 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; + m.drain_ms = (st->drain_wait_seconds - prev_drain_s) * 1000.0; + m.adopt_ms = (st->adopt_wait_seconds - prev_adopt_s) * 1000.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; + prev_drain_s = st->drain_wait_seconds; + prev_adopt_s = st->adopt_wait_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; + drain_seconds += m.drain_ms / 1000.0; + adopt_seconds += m.adopt_ms / 1000.0; } }; @@ -476,6 +486,7 @@ std::unique_ptr Session::open(const SessionConfig & cfg, im.hook->set_drop_policy(cfg.moe.drop_cold_frac, cfg.moe.drop_renorm, cfg.moe.drop_prefill); im.hook->set_predict_log(cfg.moe.predict_log); im.hook->set_predict_prefetch(cfg.moe.predict_prefetch, cfg.moe.predict_spec_max); + im.hook->set_route_ahead(cfg.moe.route_ahead); llama_context_params cparams = llama_context_default_params(); cparams.n_ctx = cfg.n_ctx; @@ -759,6 +770,7 @@ std::unique_ptr Session::open(const SessionConfig & cfg, ri.overlap = cfg.moe.enabled && cfg.moe.overlap; ri.io_two_wave = cfg.moe.enabled && cfg.moe.io_two_wave; ri.prefetch_layers = cfg.moe.enabled ? cfg.moe.prefetch_layers : 0; + ri.route_ahead = cfg.moe.enabled ? cfg.moe.route_ahead : 0; ri.predict_prefetch = cfg.moe.enabled && cfg.moe.predict_prefetch; ri.predict_log = cfg.moe.enabled && cfg.moe.predict_log; ri.predict_spec_max = cfg.moe.enabled ? cfg.moe.predict_spec_max : 0; @@ -1111,6 +1123,8 @@ RunResult Session::generate(const GenerateRequest & req, tally.prev_io_s = st0.read_seconds; tally.prev_mgmt_s = st0.mgmt_seconds; tally.prev_stall_s = st0.stall_seconds; + tally.prev_drain_s = st0.drain_wait_seconds; + tally.prev_adopt_s = st0.adopt_wait_seconds; } const IExpertSource::Stats st_spec0 = moe.enabled ? im.source.stats() : IExpertSource::Stats{}; long long prev_spec_bytes = (long long) st_spec0.spec_read_bytes; @@ -1120,6 +1134,10 @@ RunResult Session::generate(const GenerateRequest & req, // 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(); + // Per-token cursors for the hook's own eval-thread meters (route-ahead issue + watchdog): the + // hook accumulates for the session, the rows want this token's share. + long long prev_ra_issue_ns = im.hook->route_ahead_issue_ns(); + long long prev_ra_wd_ns = im.hook->route_ahead_wd_ns(); // The decode bracket below measures llama_decode and nothing else, which is what makes // compute_ms a clean residual — but it also means everything BETWEEN two decodes (sampling, @@ -1316,6 +1334,18 @@ RunResult Session::generate(const GenerateRequest & req, const double wall = secs(s0, s1); gen_seconds += wall; const IExpertSource::Stats st = moe.enabled ? im.source.stats() : IExpertSource::Stats{}; + // Route-ahead's eval-thread meters accumulate per DECODE, and one decode can confirm a whole + // group, so they are read once here and charged to the group's first row like every other + // group cost. Reading them inside the loop would advance the cursors once per token and + // credit the second and later rows with a delta of zero for the wrong reason. + double ra_issue_ms = 0.0, ra_wd_ms = 0.0; + { + const long long ri = im.hook->route_ahead_issue_ns(), rw = im.hook->route_ahead_wd_ns(); + ra_issue_ms = (ri - prev_ra_issue_ns) / 1e6; + ra_wd_ms = (rw - prev_ra_wd_ns) / 1e6; + prev_ra_issue_ns = ri; + prev_ra_wd_ns = rw; + } for (size_t e = 0; e < confirmed.size() && n_gen < req.n_predict; ++e) { const llama_token out = confirmed[e]; char piece[256]; @@ -1334,6 +1364,8 @@ RunResult Session::generate(const GenerateRequest & req, // Charged to the group's first row like every other group cost. This is a SLICE of // loop_overhead_ms, not an addition to it: both measure time outside the decode. m.mtp_draft_ms = e == 0 ? draft_s * 1000.0 : 0.0; + m.ra_issue_ms = e == 0 ? ra_issue_ms : 0.0; + m.ra_wd_ms = e == 0 ? ra_wd_ms : 0.0; m.piece = delta; // Only when someone will read it: the parser cannot resume, so this re-parses everything // generated so far on every token, and off the chat path it is a full copy of the same. @@ -1410,6 +1442,10 @@ RunResult Session::generate(const GenerateRequest & req, 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.cache_evictions = st.evictions; + s.cache_rereads = st.rereads; + s.moe_drain_s_per_token = n_gen ? tally.drain_seconds / n_gen : 0.0; + s.moe_adopt_s_per_token = n_gen ? tally.adopt_seconds / n_gen : 0.0; 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); @@ -1438,6 +1474,18 @@ RunResult Session::generate(const GenerateRequest & req, s.predict_self_by_layer = im.hook->predict_self_by_layer(); s.predict_unscored = im.hook->predict_unscored(); } + if (moe.route_ahead > 0) { + // Session totals, like the probe's: every overridden routing is an equally valid sample of + // the perturbation the flag buys, whichever turn produced it. + s.route_ahead_overridden = im.hook->route_ahead_overridden(); + s.route_ahead_passthrough = im.hook->route_ahead_passthrough(); + s.route_ahead_slots = im.hook->route_ahead_slots(); + s.route_ahead_hits = im.hook->route_ahead_hits(); + s.route_ahead_gemv_ns = im.hook->route_ahead_gemv_ns(); + s.route_ahead_gemv_jobs = im.hook->route_ahead_gemv_jobs(); + s.route_ahead_issue_ns = im.hook->route_ahead_issue_ns(); + s.route_ahead_wd_ns = im.hook->route_ahead_wd_ns(); + } if (sink) sink->on_summary(s); // One non-partial parse of the finished generation, shared by the returned text and the history diff --git a/core/src/metrics/csv_metrics_sink.cpp b/core/src/metrics/csv_metrics_sink.cpp index 6299366..c77cd6e 100644 --- a/core/src/metrics/csv_metrics_sink.cpp +++ b/core/src/metrics/csv_metrics_sink.cpp @@ -37,10 +37,10 @@ public: std::fprintf(f_, "# moe_stream=%d cache_mb=%d cache_auto=%d cache_floor_mb=%d cache_ceil_mb=%d " "force_cache=%d load_all=%d io_threads=%d o_direct=%d overlap=%d io_two_wave=%d prefetch=%d " - "predict_prefetch=%d predict_log=%d predict_spec_max=%d prefetch_sync=%d " + "route_ahead=%d predict_prefetch=%d predict_log=%d predict_spec_max=%d prefetch_sync=%d " "dense_weights=%s drop_cold_frac=%.4g drop_renorm=%d drop_prefill=%d\n", r.moe_stream, r.cache_mb, r.cache_auto, r.cache_floor_mb, r.cache_ceil_mb, r.force_cache, - r.load_all, r.io_threads, r.o_direct, r.overlap, r.io_two_wave, r.prefetch_layers, + r.load_all, r.io_threads, r.o_direct, r.overlap, r.io_two_wave, r.prefetch_layers, r.route_ahead, r.predict_prefetch, r.predict_log, r.predict_spec_max, r.prefetch_sync, r.dense_weights.c_str(), (double) r.drop_cold_frac, r.drop_renorm, r.drop_prefill); std::fprintf(f_, @@ -55,12 +55,13 @@ public: write_header(); // a caller that never sent RunInfo still gets a readable file std::fprintf(f_, "%d,%d,%.3f,%.3f,%.3f,%llu,%.2f,%.3f,%.3f,%llu,%.3f,%.3f,%d,%.2f,%.1f,%.1f,%.1f," - "%.1f,%.1f,%.1f,%.1f,%.1f,%.3f,%d,%.3f\n", + "%.1f,%.1f,%.1f,%.1f,%.1f,%.3f,%d,%.3f,%.3f,%.3f,%.3f,%.3f\n", m.step, m.steps, m.wall_ms, m.io_ms, m.compute_ms, (unsigned long long) m.read_bytes, m.cache_hit_pct, m.stall_ms, m.mgmt_ms, (unsigned long long) m.majflt, m.cpu_ms, m.dense_resident_frac, m.turn, m.majflt_mib, m.cache_budget_mib, m.rss_mib, m.rss_anon_mib, m.rss_file_mib, m.swap_mib, m.mem_available_mib, m.mem_free_mib, m.swap_free_mib, - m.loop_overhead_ms, m.mtp_batch, m.mtp_draft_ms); + m.loop_overhead_ms, m.mtp_batch, m.mtp_draft_ms, m.drain_ms, m.adopt_ms, m.ra_issue_ms, + m.ra_wd_ms); std::fflush(f_); } void on_summary(const RunSummary & s) override { @@ -75,7 +76,10 @@ public: "majflt/tok=%.2f cpu_s/tok=%.4f token_demand_MiB=%.1f layer_demand_MiB=%.1f " "experts_routed=%lld experts_dropped=%lld loop_overhead_s/tok=%.4f " "mtp_drafted=%lld mtp_accepted=%lld mtp_decodes=%lld mtp_draft_s/tok=%.4f " - "drafted_steps=%lld\n", + "drafted_steps=%lld " + "ra_committed=%lld ra_passthrough=%lld ra_agree_pct=%.1f ra_gemv_ms/tok=%.2f " + "ra_issue_ms/tok=%.2f ra_wd_ms/tok=%.2f drain_s/tok=%.3f adopt_s/tok=%.3f " + "evictions=%lld rereads=%lld\n", s.n_generated, s.s_per_token, s.tokens_per_second, s.moe_read_mib, s.moe_io_seconds, s.moe_compute_s_per_token, s.moe_io_s_per_token, s.cache_hit_pct, s.n_prompt, s.load_seconds, s.prefill_seconds, s.prefill_seconds > 0 ? s.n_prompt / s.prefill_seconds : 0.0, @@ -83,7 +87,12 @@ public: s.cache_resizes, s.moe_spec_read_mib, s.moe_spec_experts, s.moe_spec_useful, s.majflt_per_token, s.cpu_s_per_token, s.token_demand_mib, s.layer_demand_mib, s.experts_routed, s.experts_dropped, s.loop_overhead_s_per_token, s.mtp_drafted, s.mtp_accepted, s.mtp_decodes, s.mtp_draft_s_per_token, - s.drafted_steps); + s.drafted_steps, s.route_ahead_overridden, s.route_ahead_passthrough, + s.route_ahead_slots > 0 ? 100.0 * s.route_ahead_hits / s.route_ahead_slots : 0.0, + s.n_generated ? s.route_ahead_gemv_ns / 1e6 / s.n_generated : 0.0, + s.n_generated ? s.route_ahead_issue_ns / 1e6 / s.n_generated : 0.0, + s.n_generated ? s.route_ahead_wd_ns / 1e6 / s.n_generated : 0.0, s.moe_drain_s_per_token, + s.moe_adopt_s_per_token, s.cache_evictions, s.cache_rereads); std::fflush(f_); } @@ -97,10 +106,14 @@ private: // compute residual), majflt/cpu_ms (the fault + CPU-time decomposition of what remains of // "compute"), dense_resident_frac (how much of the dense set is still in RAM; -1 = unmeasured), // then the memory block. Consumers read columns by name, so order is not a contract. + // The tail names costs that would otherwise hide inside another column: mtp_batch and + // mtp_draft_ms describe a speculative group (one decode, several rows, the cost on the + // first), and drain_ms / ra_issue_ms / ra_wd_ms come out of compute_ms while adopt_ms comes + // out of mgmt_ms — see TokenMetrics. All 0 when the feature that pays them is off. std::fprintf(f_, "step,steps,wall_ms,io_ms,compute_ms,read_bytes,cache_hit_pct,stall_ms,mgmt_ms,majflt,cpu_ms," "dense_resident_frac,turn,majflt_mib,cache_budget_mib,rss_mib,rss_anon_mib," "rss_file_mib,swap_mib,mem_available_mib,mem_free_mib,swap_free_mib,loop_overhead_ms," - "mtp_batch,mtp_draft_ms\n"); + "mtp_batch,mtp_draft_ms,drain_ms,adopt_ms,ra_issue_ms,ra_wd_ms\n"); } std::FILE * f_ = nullptr; diff --git a/core/src/moe/expert_stream_source.cpp b/core/src/moe/expert_stream_source.cpp index 91ae337..d3d3ede 100644 --- a/core/src/moe/expert_stream_source.cpp +++ b/core/src/moe/expert_stream_source.cpp @@ -46,6 +46,11 @@ bool ExpertStreamSource::init(const std::vector & shard_paths, overlap_ = cfg.overlap; two_wave_ = cfg.io_two_wave; prefetch_sync_ = cfg.prefetch_sync && !cfg.overlap; // serial only: overlap lane 0 is a worker + // Under route-ahead the speculated ids of a layer ARE the ids its topk will commit, so the + // demand load adopts that layer's in-flight speculation instead of discarding and re-reading + // it. Off for every guessing predictor: adopting a wrong guess would make the load wait on + // reads it does not need. + spec_adopt_ = cfg.route_ahead > 0; cache_max_ = (size_t) std::max(0, cfg.cache_mb) * 1024ull * 1024ull; io_threads_ = std::max(1, std::min(MoeStreamConfig::io_threads_max, cfg.io_threads)); page_ = pio::vm_page(); // the real OS page size, for the dense-residency probe in any cache mode @@ -165,6 +170,7 @@ bool ExpertStreamSource::init(const std::vector & shard_paths, cnext_.assign(n_entry, -1); cspec_.assign(n_entry, 0); spec_remaining_.assign(n_entry, 0); + ever_evicted_.assign(n_entry, 0); chead_ = ctail_ = -1; cresident_ = 0; cgen_ = 0; @@ -393,14 +399,26 @@ void ExpertStreamSource::prefetch(int il, const int32_t * ids, int n_ids) { staged.clear(); staged_ids.clear(); staged_counts.clear(); - for (int i = 0; i < n_ids; ++i) { - const int e = ids[i]; - if (e < 0 || e >= n_expert_) continue; - const int32_t id = il * n_expert_ + e; - { - std::lock_guard lk(io_mtx_); + + // Filter under ONE lock, not one per expert. This runs on the eval thread while every I/O lane + // is taking the same mutex to pull its next read index, and a blocked eval thread is a + // descheduled compute thread: measured on a phone, the per-expert acquisition was part of ~120 + // acquisitions per token that cost ~600 ms of wall (the same run with a single lane, hence no + // contention, spent a quarter of the compute time on identical work). + std::vector & cand = spec_cand_; + cand.clear(); + { + std::lock_guard lk(io_mtx_); + for (int i = 0; i < n_ids; ++i) { + const int e = ids[i]; + if (e < 0 || e >= n_expert_) continue; + const int32_t id = il * n_expert_ + e; if (cvalid_[id] || spec_remaining_[id] != 0) continue; // already resident or already queued + cand.push_back(e); } + } + for (int e : cand) { + const int32_t id = il * n_expert_ + e; const size_t mark = staged.size(); int njobs = 0; bool ok = true; @@ -426,6 +444,7 @@ void ExpertStreamSource::prefetch(int il, const int32_t * ids, int n_ids) { break; } if (njobs == 0) continue; + if (ever_evicted_[id]) ++rereads_; // same accounting on the speculative path staged_ids.push_back(id); staged_counts.push_back(njobs); any = true; @@ -451,7 +470,10 @@ void ExpertStreamSource::prefetch(int il, const int32_t * ids, int n_ids) { const bool ok = read_slice(0, j); if (ok && g == spec_gen_) { spec_read_bytes_.fetch_add((long long) j.nbytes); - if (--spec_remaining_[j.flag] == 0) spec_done_.push_back(j.flag); + if (--spec_remaining_[j.flag] == 0) { + spec_done_.push_back(j.flag); + spec_done_pending_.fetch_add(1, std::memory_order_relaxed); + } } --spec_inflight_; } @@ -480,7 +502,17 @@ void ExpertStreamSource::drain_spec(int lane, uint64_t worker_seen) { std::lock_guard lk(io_mtx_); if (ok && g == spec_gen_) { // ignore reads from a cancelled round spec_read_bytes_.fetch_add((long long) j.nbytes); - if (--spec_remaining_[j.flag] == 0) spec_done_.push_back(j.flag); + if (spec_remaining_[j.flag] > 0 && --spec_remaining_[j.flag] == 0) { + spec_done_.push_back(j.flag); + spec_done_pending_.fetch_add(1, std::memory_order_relaxed); + // The adoption wait watches per-entry counters, not just the in-flight count. + io_cv_done_.notify_all(); + } + } else if (!ok && spec_adopt_) { + // Forget-on-failure: zeroing the counter keeps the adoption wait from parking + // forever, and !cvalid means a later demand read re-buys the entry. + spec_remaining_[j.flag] = 0; + io_cv_done_.notify_all(); } if (--spec_inflight_ == 0) io_cv_done_.notify_all(); } @@ -490,7 +522,72 @@ void ExpertStreamSource::drain_spec(int lane, uint64_t worker_seen) { // Quiesce speculation before real staging: cancel queued reads, wait out in-flight ones, then on // this (eval) thread integrate every fully-read entry into the cache and release the rest. All LRU // mutation happens here, single-threaded, so it never races the real staging that follows. -void ExpertStreamSource::quiesce_spec() { +void ExpertStreamSource::quiesce_spec(int adopt_il) { + // Adoption (route-ahead only): the layer about to stage COMMITTED to the ids its speculation + // was issued for, so its queued spec reads are not bets to discard — they are the layer's own + // demand reads, already staged and possibly partly done. Keep exactly those, drop the rest of + // the queue, and FINISH them here (this thread drains alongside any idle lane) before the + // normal quiesce below integrates them as resident entries — which the staging then sees as + // hits, so no byte is read twice. Without this, a committed layer's in-flight speculation was + // released at its own load and re-read on demand: measured on the host at depth 2, every + // early read was wasted that way (0% useful, double flash per token). The generation is NOT + // bumped before the drain — bumping is what disowns reads, and these are being adopted. + // Adoption (route-ahead only): under a committed routing, EVERY job in the spec queue is a + // read some layer will demand verbatim — the loading layer's own jobs are its demand reads + // already staged, and the other layers' jobs are demand reads a few callbacks early. So + // nothing is cancelled: the loading layer's jobs move to the front and are finished now, + // completed entries integrate, and everything else stays queued for the idle lanes. The + // destructive quiesce below (cancel, disown, release) remains the right treatment for the + // guessing predictors, whose queue really is bets. Measured before this branch existed: at + // depth 2 every early read was destroyed — by the intervening load's quiesce or by the + // settle preceding each issue — for 0% useful and double flash per token. + if (spec_adopt_ && adopt_il >= 0) { + std::vector & adopt = spec_adopt_ids_; + adopt.clear(); + { + std::lock_guard lk(io_mtx_); + std::stable_partition(spec_jobs_.begin() + (ptrdiff_t) spec_next_, spec_jobs_.end(), + [&](const IoJob & j) { return j.layer == (int16_t) adopt_il; }); + for (int32_t id : spec_touched_) + if (id / n_expert_ == adopt_il && spec_remaining_[id] != 0) adopt.push_back(id); + // Without the full quiesce ever running, spec_touched_ would grow for the whole + // session; entries fully settled need no release sweep, so drop them here. + spec_touched_.erase(std::remove_if(spec_touched_.begin(), spec_touched_.end(), + [&](int32_t id) { return spec_remaining_[id] == 0; }), + spec_touched_.end()); + } + if (!adopt.empty()) { + io_cv_.notify_all(); // idle worker lanes take the queue + // Serially, lane 0 belongs to this thread, so it reads its own layer's jobs alongside + // the workers; under overlap every lane is a worker and the eval thread must not + // touch a reader — doing so raced worker 0 on one reader and fed the matmul + // corrupted slices (first symptom: the routing agreeing BETTER with its own + // prediction, garbage states drifting toward whatever the perturbed layers produce). + // + // Then WAIT for the loading layer's own entries, in both modes — measured, not + // assumed: a no-wait variant that skipped these entries at staging and published + // per-expert readiness from the completing reads was built and benchmarked, and it + // LOST (5.8 vs 6.6 tok/s on the host at depth 2). Late integration enters the + // entries cold, the eviction churns them, and the per-expert stalls it trades this + // wait for cost more than the wait — which is short by construction, because these + // jobs started a layer or more ago. + if (!overlap_) drain_adopted(adopt_il); + const auto tw0 = clock_t_::now(); + { + std::unique_lock lk(io_mtx_); + io_cv_done_.wait(lk, [&] { + if (io_stop_) return true; + for (int32_t id : adopt) + if (spec_remaining_[id] != 0) return false; + return true; + }); + } + adopt_wait_ns_ += + (long long) std::chrono::duration_cast(clock_t_::now() - tw0).count(); + } + spec_integrate_done(); + return; + } std::vector done, touched; { std::unique_lock lk(io_mtx_); @@ -499,6 +596,7 @@ void ExpertStreamSource::quiesce_spec() { spec_next_ = 0; io_cv_done_.wait(lk, [&] { return spec_inflight_ == 0 || io_stop_; }); done.swap(spec_done_); + spec_done_pending_.store(0, std::memory_order_relaxed); touched.swap(spec_touched_); } // Integrate completed entries (all projections resident) into the LRU cache. @@ -629,6 +727,8 @@ void ExpertStreamSource::release_entry_pages(int32_t id) { } void ExpertStreamSource::evict_tail() { const int32_t id = ctail_; + ++evictions_; + ever_evicted_[id] = 1; // the next read of this entry is the cache paying twice release_entry_pages(id); cvalid_[id] = 0; cspec_[id] = 0; @@ -658,7 +758,92 @@ void ExpertStreamSource::settle_spec() { // makes the one inside the load_layer that follows a cheap no-op: nothing left queued, and // nothing in flight. The cost is that its time lands outside the mgmt_ns_ window — which is // why a traced run is not a benchmark run. - if (active_ && cache_max_) quiesce_spec(); + // + // Under route-ahead the destructive settle would be a saboteur, not a bookkeeper: the hook + // settles right before every issue, so the full quiesce here cancelled the future layers' + // committed reads one callback before their load could adopt them. Integrate-only instead. + if (!active_ || cache_max_ == 0) return; + if (spec_adopt_) { + // Cheap exit without touching the lanes' mutex: this runs before every route-ahead issue + // (40 times a token) and usually has nothing to integrate. spec_done_pending_ is the + // publisher's own count, so a relaxed read either sees work — and we take the lock — or + // races a completion we will pick up at the next issue microseconds later. + if (spec_done_pending_.load(std::memory_order_relaxed) != 0) spec_integrate_done(); + } else { + quiesce_spec(); + } +} + +// Integrate every completed speculative entry into the LRU cache, cancelling nothing — the +// route-ahead settle. Eval-thread only (LRU mutation), like the quiesce that subsumes it. +void ExpertStreamSource::spec_integrate_done() { + std::vector & done = spec_done_scratch_; + done.clear(); + { + // BOTH eligibility guards run under io_mtx_, exactly as the destructive integrate runs + // its own, and the remaining check is load-bearing: a completion can land in spec_done_ + // concurrently with (and just before) an issue that re-stages the same entry, and + // integrating that stale completion would push an entry with LIVE queued jobs into the + // LRU — the eviction then decommits pages a lane is about to write, which was this + // feature's one segfault. The stale id is simply dropped; the new round's own completion + // re-pushes it. After this snapshot the check cannot rot: re-staging happens only on + // this same (eval) thread. + std::lock_guard lk(io_mtx_); + for (int32_t id : spec_done_) + if (!cvalid_[id] && spec_remaining_[id] == 0) done.push_back(id); + spec_done_.clear(); + spec_done_pending_.store(0, std::memory_order_relaxed); + } + for (int32_t id : done) { + cvalid_[id] = 1; + cspec_[id] = 1; // speculative until a real lookup hits it (then counted useful) + cstamp_[id] = 0; // not used this generation → evictable if the budget is tight + cresident_ += entry_bytes(id / n_expert_); + // Cold end for a GUESS (a mispredicted expert should be the first thing reclaimed) — but + // under route-ahead there are no guesses: every integrated entry is a read some layer + // COMMITTED to, at most N layers from being routed. Entering it cold handed it to the + // very evictions of the intervening layers' loads, and on a device whose cache runs at + // capacity that destroyed the entries before use and re-bought them on demand — a + // perfect prefetch paying its bytes twice (measured: +34% read/token, drop off). A + // committed read enters HOT; its own routing arrives to promote it within N layers + // anyway. + if (spec_adopt_) + lru_push_front(id); + else + lru_push_back(id); + spec_experts_.fetch_add(1); + } +} + +// Serial-mode helper for adoption: this thread owns lane 0, so it reads the loading layer's own +// adopted jobs — which sit at the front after the partition — and stops at the first job that +// belongs to a future layer, which stays queued for the worker lanes' idle time. +void ExpertStreamSource::drain_adopted(int adopt_il) { + for (;;) { + IoJob j; + { + std::lock_guard lk(io_mtx_); + if (spec_next_ >= spec_jobs_.size() || spec_jobs_[spec_next_].layer != (int16_t) adopt_il) return; + j = spec_jobs_[spec_next_++]; + ++spec_inflight_; + } + const bool ok = read_slice(0, j); + { + std::lock_guard lk(io_mtx_); + if (ok) { + spec_read_bytes_.fetch_add((long long) j.nbytes); + if (spec_remaining_[j.flag] > 0 && --spec_remaining_[j.flag] == 0) { + spec_done_.push_back(j.flag); + spec_done_pending_.fetch_add(1, std::memory_order_relaxed); + } + } else { + // A failed read forgets the entry: zeroing its counter keeps the adoption wait + // from hanging, and !cvalid means the demand path simply re-reads it. + spec_remaining_[j.flag] = 0; + } + if (--spec_inflight_ == 0) io_cv_done_.notify_all(); + } + } } void ExpertStreamSource::query_residency(int il, const int32_t * ids, int n_ids, uint8_t * out) const { @@ -723,6 +908,7 @@ bool ExpertStreamSource::touch_entry(int il, int e, bool & hit, bool promote, in if (commit_only_proj >= 0 && p != commit_only_proj) continue; if (!commit_proj_pages(il, e, p)) return false; } + if (ever_evicted_[id]) ++rereads_; // this entry was resident once; the cache is buying it again cvalid_[id] = 1; cspec_[id] = 0; // a real read, not speculative cresident_ += entry_bytes(il); @@ -746,7 +932,7 @@ bool ExpertStreamSource::load_layer(int il, const int32_t * ids, int n_ids) { // Integrate / discard any speculative prefetch before this layer's real staging touches the // cache. After this returns no spec read is in flight and completed ones are resident hits. The // budget is fixed for the run (auto sizes it once at init), so there is nothing to resize here. - if (cache_max_) quiesce_spec(); + if (cache_max_) quiesce_spec(spec_adopt_ ? il : -1); auto stage = [&](int e) -> bool { if (cache_max_ == 0) { @@ -861,8 +1047,11 @@ bool ExpertStreamSource::load_layer_async(int il, const int32_t * ids, int n_ids // safety guarantee (no worker still reading an about-to-be-evicted page) and it covers // load_all, whose surplus jobs no hook ever waits on. Serial's epilogue waited here too. { + const auto tw0 = clock_t_::now(); std::unique_lock lk(io_mtx_); io_cv_done_.wait(lk, [&] { return done_cnt_ == batch_njobs_ || io_stop_; }); + drain_wait_ns_ += + (long long) std::chrono::duration_cast(clock_t_::now() - tw0).count(); } if (fatal_.load(std::memory_order_acquire)) return false; @@ -874,7 +1063,7 @@ bool ExpertStreamSource::load_layer_async(int il, const int32_t * ids, int n_ids // Integrate / discard speculative prefetch before this layer's real staging (the previous // batch is already drained above, so this only waits on in-flight spec reads). The budget is // fixed for the run, so there is nothing to resize here. - if (cache_max_) quiesce_spec(); + if (cache_max_) quiesce_spec(spec_adopt_ ? il : -1); // 2. New generation for this layer. A flag counts as ready only once its gen matches. const uint32_t gen = async_gen_.fetch_add(1, std::memory_order_relaxed) + 1; @@ -1011,14 +1200,22 @@ bool ExpertStreamSource::load_layer_async(int il, const int32_t * ids, int n_ids // Promote every touched expert in raw id order (token-major) so the LRU order reflects // the LAST token that used it, not the sorted/first-touch order staged_ imposes — this // keeps the prompt tail's experts hot across prefill. This loop is the ONLY promotion on - // this path (touch_entry above runs with promote=false); every id staged above is valid - // and linked, a miss by its initial push, so unlink here is always legal. + // this path (touch_entry above runs with promote=false). // Skipped in load_all (everything is resident, so LRU order is meaningless). + // + // The cvalid guard is load-bearing, not defensive: an ADOPTED-PENDING expert (route-ahead) + // was deliberately never staged, so it is not in the LRU — and lru_unlink on a node whose + // prev/next are both -1 rewrites chead_/ctail_ to -1, orphaning every linked entry in one + // call. That was this feature's second segfault: the next push_front made the pending + // entry the LIST, eviction took it (decommitting pages its reads were still filling), and + // the orphaned entries could never be evicted again. A pending expert needs no promotion: + // it enters the LRU at integration, and its first real hit promotes it then. if (!load_all_) { for (int i = 0; i < n_ids; ++i) { const int e = ids[i]; if (e < 0 || e >= n_expert_) continue; const int32_t id = il * n_expert_ + e; + if (!cvalid_[id]) continue; lru_unlink(id); lru_push_front(id); } @@ -1155,6 +1352,10 @@ IExpertSource::Stats ExpertStreamSource::stats() const { s.stall_seconds = stall_ns_.load() / 1e9; s.cache_budget_bytes = (uint64_t) cache_max_; s.cache_resizes = cache_resizes_; + s.evictions = evictions_; + s.rereads = rereads_; + s.drain_wait_seconds = drain_wait_ns_ / 1e9; + s.adopt_wait_seconds = adopt_wait_ns_ / 1e9; s.dense_resident_frac = dense_.resident_frac(); s.token_demand_bytes = (uint64_t) token_demand_; s.layer_demand_bytes = (uint64_t) layer_demand_; @@ -1190,6 +1391,7 @@ void ExpertStreamSource::shutdown() { if (t.joinable()) t.join(); io_pool_.clear(); spec_done_.clear(); + spec_done_pending_.store(0, std::memory_order_relaxed); spec_touched_.clear(); for (int p = 0; p < MoeRecipe::max_exps; ++p) diff --git a/core/src/moe/expert_stream_source.h b/core/src/moe/expert_stream_source.h index e997068..90ac78b 100644 --- a/core/src/moe/expert_stream_source.h +++ b/core/src/moe/expert_stream_source.h @@ -142,7 +142,12 @@ private: // Speculative prefetch (temporal): drain queued spec reads on an idle lane, and integrate / // discard them on the eval thread at the next real load. See docs/prefetch.md. void drain_spec(int lane, uint64_t worker_seen); - void quiesce_spec(); + // adopt_il >= 0 (route-ahead only): the layer being staged COMMITTED to its speculated ids, so + // its queued spec reads are the demand reads — finish them instead of cancelling them, and + // leave every other layer's committed reads queued rather than destroying them. + void quiesce_spec(int adopt_il = -1); + void spec_integrate_done(); // integrate completed spec entries; cancel nothing (route-ahead) + void drain_adopted(int adopt_il); // serial: lane 0 reads the loading layer's adopted jobs inline void release_entry_pages(int32_t id); // Accumulate one token's routed working set. Eval-thread only, called per layer load. @@ -190,8 +195,27 @@ private: bool active_ = false; bool load_all_ = false; bool overlap_ = false; - bool two_wave_ = false; // publish the first projection's jobs before committing the rest (#118) - bool prefetch_sync_ = false; // test only: drain prefetch reads synchronously (serial mode) + bool two_wave_ = false; // publish the first projection's jobs before committing the rest (#118) + bool prefetch_sync_ = false; // test only: drain prefetch reads synchronously (serial mode) + bool spec_adopt_ = false; // route-ahead: staged layers adopt committed spec reads (see quiesce_spec) + std::vector spec_adopt_ids_; // scratch: the loading layer's pending spec entries + std::vector spec_done_scratch_; // scratch: completed entries being integrated + std::vector spec_cand_; // scratch: prefetch candidates surviving the residency filter + // How many completed entries are waiting to be integrated. Mirrors spec_done_.size(), but + // readable without the lanes' mutex, so the route-ahead settle — which runs before every issue, + // 40 times a token — can skip the lock entirely on the common empty case. + std::atomic spec_done_pending_{0}; + // Churn detector: which entries have ever been evicted, and how many reads went to an entry + // that had been resident before. A prefetch cannot reduce the bytes a routing needs — the + // ideal is the same bytes, earlier — so any read of something the cache already had once is + // the cache paying twice, and it is the only way a "100% useful" prefetch can still raise + // the byte count. Reported per run; cheap enough (one byte per entry, one branch per read). + std::vector ever_evicted_; + long long rereads_ = 0, evictions_ = 0; + // Eval-thread wait counters (see IExpertSource::Stats): the async load's previous-batch drain + // (lives in the compute residual) and the route-ahead adoption wait (lives inside mgmt). + // Eval-thread only, like the LRU itself — two clock reads per wait, nothing on the lanes. + long long drain_wait_ns_ = 0, adopt_wait_ns_ = 0; int n_layer_ = 0; int n_expert_ = 0; size_t align_ = 4096; diff --git a/core/src/moe/router_hook.cpp b/core/src/moe/router_hook.cpp index 248c15c..f49007b 100644 --- a/core/src/moe/router_hook.cpp +++ b/core/src/moe/router_hook.cpp @@ -230,6 +230,7 @@ void RouterHook::apply_drop(ggml_tensor * wt) { source_->load_layer(D.layer, drop_ids_.data(), (int) drop_ids_.size()); D.deferred = false; predict_after_load(D.layer); + route_ahead_collect(D.layer); } // Finish with the layer whose topk we last saw: record which node ended its weight chain, so the @@ -250,6 +251,7 @@ void RouterHook::close_drop_layer() { if (D.layer < (int) term_variant_.size()) term_variant_[D.layer] = -1; D.deferred = false; predict_after_load(D.layer); + route_ahead_collect(D.layer); } D.layer = -1; } @@ -274,6 +276,27 @@ void RouterHook::set_predict_prefetch(bool on, int spec_max) { if (!on) predict_worker_stop(); } +void RouterHook::set_route_ahead(int n) { + route_ahead_ = n > 0 ? n : 0; + const int nl = n_layer_ > 0 ? n_layer_ : 0; + ra_pred_.assign((size_t) nl, std::vector{}); + ra_gate_q_.assign((size_t) nl, std::vector{}); + ra_gate_s_.assign((size_t) nl, std::vector{}); + ra_overridden_ = ra_passthrough_ = ra_slots_ = ra_hits_ = 0; + ra_tripped_ = false; + ra_gemv_ns_.store(0); + ra_gemv_jobs_.store(0); + ra_issue_ns_ = 0; + // The pointer stashes are normally sized by predict_reset(); route-ahead can be armed with + // both probe and prefetch off, so make sure they exist rather than assume who ran first. + if (gate_w_.size() != (size_t) nl) gate_w_.assign((size_t) nl, nullptr); + if (h_t_.size() != (size_t) nl) h_t_.assign((size_t) nl, nullptr); + // The GEMV runs on the shared prediction worker (predict_prefetch is mutually exclusive, so + // the job slot is never contended); own the thread's lifecycle the same way it does. + if (route_ahead_ > 0 && !pred_worker_.joinable()) pred_worker_ = std::thread([this] { predict_worker_main(); }); + if (route_ahead_ == 0 && !predict_prefetch_) predict_worker_stop(); +} + RouterHook::~RouterHook() { predict_worker_stop(); } @@ -324,11 +347,70 @@ void RouterHook::predict_worker_main() { pred_job_pending_ = false; } std::vector scores; - if (!job.gate || !gate_scores(job.gate, job.row, scores)) continue; + // Time the GEMV itself (the ranking that follows is O(k*n_expert) and rides along): this + // is the feature's own CPU bill, invisible in wall time wherever a spare core exists. + const auto tg0 = std::chrono::steady_clock::now(); + // Commit jobs read the int8 mirror — a quarter of the bytes, which on a phone IS the cost: + // the float GEMV there is bound by pulling ~4 MB of gate matrix per layer across a memory + // bus the decode is already saturating. The probe's jobs always keep the exact weights, + // because its whole purpose is to be a control. The mirror is built here, on this worker, + // the first time a layer is predicted for. + // + // aarch64 only, and measured both ways: on x86 the float path is fully vectorized and DRAM + // bandwidth is ample, so the int8 detour costs more than it saves (0.38 vs 0.26 ms per + // GEMV on the host, same config). The bytes only matter where they are scarce. + bool scored = false; +#if defined(__aarch64__) + if (job.commit && job.nl >= 0 && quantize_gate(job.nl)) scored = gate_scores_q(job.nl, job.row, scores); +#endif + if (!scored) scored = job.gate && gate_scores(job.gate, job.row, scores); + if (job.commit) { + ra_gemv_ns_.fetch_add( + (long long) std::chrono::duration_cast(std::chrono::steady_clock::now() - tg0) + .count()); + ra_gemv_jobs_.fetch_add(1); + } + if (!scored) continue; PredictResult r; r.seq = job.seq; r.nl = job.nl; - build_spec_lists(scores, job.nu, job.drop_frac, job.spec_max, job.resident, r.spec, r.keep); + r.commit = job.commit; + if (job.commit) { + // Rank exactly the routing's width, not the probe's max: rank_top_k is a partial + // selection whose cost grows with the square of the width, so ranking 32 to use 8 + // was four times the work for entries nothing reads — measurable on a phone, where + // this runs 40 times a token beside a decode competing for the same cores. + rank_top_k(scores, job.nu, r.ids); + // The issue list rides in r.spec: of the top-nu the commit will route, only the + // experts the drop policy would actually READ. Same arithmetic as build_spec_lists + // (softmax over the top-nu as a proxy for the routing weights, threshold at + // drop_frac/nu, the top expert always kept) — an expert predicted below the drop + // threshold is left cold ON PURPOSE, so the commit-time drop discards it unread + // exactly as it does the router's own cold tail. Measured on device before this + // filter existed: uncapped early reads made every committed expert resident, the + // residency-gated drop stopped biting at all, and a run paid 3x the flash and half + // the tok/s for a perfect prefetch of bytes the baseline never read (78 -> 244 + // MB/token at depth 4, drop 0.75). With drop off the threshold is 0 and the list + // is simply the committed top-nu. + const int nu = job.nu < (int) r.ids.size() ? job.nu : (int) r.ids.size(); + if (nu > 0) { + float w[predict_max_k]; + float mx = scores[(size_t) r.ids[0]]; + for (int k = 1; k < nu; ++k) + if (scores[(size_t) r.ids[k]] > mx) mx = scores[(size_t) r.ids[k]]; + float sum = 0.0f; + for (int k = 0; k < nu; ++k) { + w[k] = std::exp(scores[(size_t) r.ids[k]] - mx); + sum += w[k]; + } + const float thr = job.drop_frac > 0.0f ? job.drop_frac / (float) nu : 0.0f; + for (int k = 0; k < nu; ++k) + if (k == 0 || job.drop_frac <= 0.0f || (sum > 0.0f && w[k] / sum >= thr)) + r.spec.push_back(r.ids[(size_t) k]); + } + } else { + build_spec_lists(scores, job.nu, job.drop_frac, job.spec_max, job.resident, r.spec, r.keep); + } r.ready = true; std::lock_guard lk(pred_mtx_); pred_result_ = std::move(r); @@ -365,30 +447,63 @@ static bool row_to_float(const ggml_tensor * t, int j, std::vector & out) // per token on a 40-layer, 256-expert model, which measured as ~35-45 ms per GEMV pass and // dwarfed everything else this feature does. On aarch64 the compiler converts __fp16 natively // (one instruction, vectorizable), so that path is used wherever it exists. +// The four-accumulator shape is not a micro-optimisation, it is what makes this loop vectorize at +// all: floating-point addition is not associative, so a compiler may not re-order a single- +// accumulator reduction, and the obvious `acc += p[d] * h[d]` compiles to one scalar multiply-add +// per cycle however high the optimisation level. Splitting the reduction into independent partial +// sums gives the vectorizer permission it cannot take on its own. Measured on the host: 0.77 -> +// 0.20 ms per gate pass on a 256-expert model — and this is the whole feature's per-layer tax, so +// on a 4-thread phone (where it competes with the decode for cores AND for DRAM bandwidth) it was +// the difference between a win and a rout. The partial sums change the summation order, so the +// last bits of a score may differ from the graph's own gate; the ranking is unaffected in every +// case that is not already a numerical tie, and the watchdog is what would catch it if it were. static bool gate_scores(const ggml_tensor * w, const std::vector & h, std::vector & out) { const int64_t nd = w->ne[0], ne = w->ne[1]; if (nd != (int64_t) h.size() || ne <= 0) return false; if (w->type != GGML_TYPE_F32 && w->type != GGML_TYPE_F16) return false; out.assign((size_t) ne, 0.0f); + const float * hp = h.data(); + const int64_t nd4 = nd & ~(int64_t) 3; for (int64_t e = 0; e < ne; ++e) { const char * row = (const char *) w->data + (size_t) e * w->nb[1]; - float acc = 0.0f; + float a0 = 0.0f, a1 = 0.0f, a2 = 0.0f, a3 = 0.0f; if (w->type == GGML_TYPE_F32) { const float * p = (const float *) row; - for (int64_t d = 0; d < nd; ++d) - acc += p[d] * h[(size_t) d]; + for (int64_t d = 0; d < nd4; d += 4) { + a0 += p[d + 0] * hp[d + 0]; + a1 += p[d + 1] * hp[d + 1]; + a2 += p[d + 2] * hp[d + 2]; + a3 += p[d + 3] * hp[d + 3]; + } + for (int64_t d = nd4; d < nd; ++d) + a0 += p[d] * hp[d]; } else { #if defined(__aarch64__) const __fp16 * p = (const __fp16 *) row; - for (int64_t d = 0; d < nd; ++d) - acc += (float) p[d] * h[(size_t) d]; #else - const ggml_fp16_t * p = (const ggml_fp16_t *) row; - for (int64_t d = 0; d < nd; ++d) - acc += ggml_fp16_to_fp32(p[d]) * h[(size_t) d]; + const ggml_fp16_t * q = (const ggml_fp16_t *) row; +#endif + for (int64_t d = 0; d < nd4; d += 4) { +#if defined(__aarch64__) + a0 += (float) p[d + 0] * hp[d + 0]; + a1 += (float) p[d + 1] * hp[d + 1]; + a2 += (float) p[d + 2] * hp[d + 2]; + a3 += (float) p[d + 3] * hp[d + 3]; +#else + a0 += ggml_fp16_to_fp32(q[d + 0]) * hp[d + 0]; + a1 += ggml_fp16_to_fp32(q[d + 1]) * hp[d + 1]; + a2 += ggml_fp16_to_fp32(q[d + 2]) * hp[d + 2]; + a3 += ggml_fp16_to_fp32(q[d + 3]) * hp[d + 3]; +#endif + } + for (int64_t d = nd4; d < nd; ++d) +#if defined(__aarch64__) + a0 += (float) p[d] * hp[d]; +#else + a0 += ggml_fp16_to_fp32(q[d]) * hp[d]; #endif } - out[(size_t) e] = acc; + out[(size_t) e] = (a0 + a1) + (a2 + a3); } return true; } @@ -423,6 +538,277 @@ void RouterHook::rank_top_k(const std::vector & scores, int k, std::vecto } } +// Build the int8 mirror of layer il's gate matrix (see the header for why). One symmetric scale +// per expert ROW keeps each row's dynamic range: the rows are what get compared against one +// another, and a single matrix-wide scale would flatten a quiet row into noise. Runs once per +// layer, on the prediction worker, reading a weight leaf nothing mutates. +bool RouterHook::quantize_gate(int il) { + if (il < 0 || il >= (int) ra_gate_q_.size()) return false; + if (!ra_gate_q_[il].empty()) return true; // already mirrored + const ggml_tensor * w = gate_w_[il]; + if (!w || !w->data) return false; + if (w->type != GGML_TYPE_F32 && w->type != GGML_TYPE_F16) return false; + const int64_t nd = w->ne[0], ne = w->ne[1]; + if (nd <= 0 || ne <= 0) return false; + + std::vector q((size_t) nd * (size_t) ne); + std::vector s((size_t) ne, 0.0f); + std::vector row((size_t) nd); + for (int64_t e = 0; e < ne; ++e) { + const char * src = (const char *) w->data + (size_t) e * w->nb[1]; + float amax = 0.0f; + for (int64_t d = 0; d < nd; ++d) { + float v; + if (w->type == GGML_TYPE_F32) { + v = ((const float *) src)[d]; + } else { +#if defined(__aarch64__) + v = (float) ((const __fp16 *) src)[d]; +#else + v = ggml_fp16_to_fp32(((const ggml_fp16_t *) src)[d]); +#endif + } + row[(size_t) d] = v; + const float a = v < 0.0f ? -v : v; + if (a > amax) amax = a; + } + const float scale = amax > 0.0f ? amax / 127.0f : 0.0f; + s[(size_t) e] = scale; + const float inv = scale > 0.0f ? 1.0f / scale : 0.0f; + int8_t * dst = q.data() + (size_t) e * (size_t) nd; + for (int64_t d = 0; d < nd; ++d) { + const float t = row[(size_t) d] * inv; + const int v = (int) (t < 0.0f ? t - 0.5f : t + 0.5f); + dst[d] = (int8_t) (v < -127 ? -127 : (v > 127 ? 127 : v)); + } + } + ra_gate_q_[il] = std::move(q); + ra_gate_s_[il] = std::move(s); + return true; +} + +// The prediction's GEMV against the int8 mirror. The ACTIVATION is quantized too, once per call, +// so the inner loop is int8 x int8 accumulated in int32 — the shape ARM's dotprod instructions +// exist for (this build targets armv8.2-a+dotprod), and the shape a compiler will vectorize +// without needing permission to reorder a float reduction. Both effects point the same way on a +// phone: a quarter of the bytes off the memory bus, and integer MACs instead of float ones. +// Ranking only needs the ORDER of the scores, and both scales are positive, so the per-expert +// scale is applied at the end where it costs one multiply per expert instead of one per element. +bool RouterHook::gate_scores_q(int il, const std::vector & h, std::vector & out) { + if (il < 0 || il >= (int) ra_gate_q_.size() || ra_gate_q_[il].empty()) return false; + const ggml_tensor * w = gate_w_[il]; + if (!w) return false; + const int64_t nd = w->ne[0], ne = w->ne[1]; + if (nd != (int64_t) h.size() || (size_t) nd * (size_t) ne != ra_gate_q_[il].size()) return false; + + // Quantize the activation row: one pass over n_embd, amortised across all n_expert rows. + float amax = 0.0f; + for (int64_t d = 0; d < nd; ++d) { + const float a = h[(size_t) d] < 0.0f ? -h[(size_t) d] : h[(size_t) d]; + if (a > amax) amax = a; + } + if (!(amax > 0.0f)) return false; // a zero row says nothing; let the caller fall back + const float hs = amax / 127.0f; + std::vector & qh = ra_qh_; + qh.resize((size_t) nd); + { + const float inv = 1.0f / hs; + for (int64_t d = 0; d < nd; ++d) { + const float t = h[(size_t) d] * inv; + const int v = (int) (t < 0.0f ? t - 0.5f : t + 0.5f); + qh[(size_t) d] = (int8_t) (v < -127 ? -127 : (v > 127 ? 127 : v)); + } + } + + const int8_t * q = ra_gate_q_[il].data(); + const float * sc = ra_gate_s_[il].data(); + const int8_t * hp = qh.data(); + const int64_t nd4 = nd & ~(int64_t) 3; + out.assign((size_t) ne, 0.0f); + for (int64_t e = 0; e < ne; ++e) { + const int8_t * p = q + (size_t) e * (size_t) nd; + int32_t a0 = 0, a1 = 0, a2 = 0, a3 = 0; + for (int64_t d = 0; d < nd4; d += 4) { + a0 += (int32_t) p[d + 0] * (int32_t) hp[d + 0]; + a1 += (int32_t) p[d + 1] * (int32_t) hp[d + 1]; + a2 += (int32_t) p[d + 2] * (int32_t) hp[d + 2]; + a3 += (int32_t) p[d + 3] * (int32_t) hp[d + 3]; + } + for (int64_t d = nd4; d < nd; ++d) + a0 += (int32_t) p[d] * (int32_t) hp[d]; + out[(size_t) e] = (float) ((a0 + a1) + (a2 + a3)) * sc[e] * hs; + } + return true; +} + +// At the topk of layer il, BEFORE the commit overwrites the ids: sample the watchdog against the +// router's own choice, copy the gate-input row, and submit the ranking job for layer il+N to the +// worker. This is the prefetch's barrier-less read (the row was stashed in the ask pass and is +// only PROBABLY still intact by now — ggml's planner may reuse the buffer), made safe for a +// committed consumer by two things: the sampled fresh-gate control below, which disarms the whole +// policy out loud if the read stops reproducing the router; and the passthrough default, which +// keeps any unproven layer on the router's own routing. +void RouterHook::route_ahead_submit(ggml_tensor * ids, int il, int nu, int nt) { + if (ra_tripped_ || il < 0 || il >= n_layer_ || nu <= 0 || nu > predict_max_k || nt <= 0) return; + // Invalidate this token's slot for the target layer FIRST, before any early return below can + // skip it: what sits there is the PREVIOUS token's ranking for that layer, and a submit that + // declines (an unstashed gate, an unreadable row) would otherwise leave it in place for + // apply_route_ahead to commit — routing a token by a prediction made from another token's + // hidden state. Passthrough is the correct behaviour for a layer this token cannot predict. + const int target = il + route_ahead_; + if (target < n_layer_) ra_pred_[target].clear(); + ggml_tensor * h = h_t_[il]; + if (!h || !h->data || h->ne[1] <= 0) return; + const int j = (int) h->ne[1] - 1; // the batch's last token, the row every predictor here uses + + // Watchdog: the layer's OWN gate on the just-copied row must reproduce the ids the router just + // chose — read from the tensor now, before apply_route_ahead rewrites them. Same cadence and + // thresholds as the prefetch's; the difference is what a trip protects against. + if (++wd_rout_ % wd_interval == 0) { + // The sampled control is an exact float GEMV on the EVAL thread — cheap because it is + // sampled, but it lives in the compute residual, so it carries its own meter. + const auto tw0 = std::chrono::steady_clock::now(); + if (gate_w_[il] && row_to_float(h, j, pred_row_) && gate_scores(gate_w_[il], pred_row_, pred_scores_)) { + std::vector ctrl; + rank_top_k(pred_scores_, nu, ctrl); + int hits = 0; + for (int k = 0; k < nu; ++k) { + const int32_t actual = *id_at(ids, nt - 1, k); + for (int32_t p : ctrl) + if (p == actual) { + ++hits; + break; + } + } + wd_slots_ += nu; + wd_hits_ += hits; + if (wd_slots_ >= wd_min_slots && (double) wd_hits_ < wd_min_frac * (double) wd_slots_) { + ra_tripped_ = true; + std::fprintf(stderr, + "bmoe: route-ahead disarmed — control at %.1f%% over %lld slots says the " + "barrier-less gate-input read is not trustworthy on this graph; routing " + "stays the router's own from here\n", + 100.0 * wd_hits_ / wd_slots_, wd_slots_); + } + } + ra_wd_ns_ += + (long long) std::chrono::duration_cast(std::chrono::steady_clock::now() - tw0) + .count(); + if (ra_tripped_) return; + } + + const int nl = target; + if (nl >= n_layer_) return; + const ggml_tensor * wn = gate_w_[nl]; + if (!wn || !wn->data) return; // target matrix not stashed yet (first decode token): passthrough + if (!row_to_float(h, j, pred_row_)) return; + { + std::lock_guard lk(pred_mtx_); + pred_job_.seq = ++pred_seq_; + pred_job_.nl = nl; + pred_job_.nu = nu; + pred_job_.commit = true; + pred_job_.drop_frac = drop_frac_; // the issue list is drop-aware; the commit is not + pred_job_.gate = wn; + pred_job_.row = pred_row_; + pred_job_pending_ = true; + } + pred_cv_.notify_one(); +} + +// Pop a finished commit ranking, right after layer il's load was handed to the source — the same +// post-load window every speculation issues in. A result for a layer still ahead of il gets its +// reads started HERE: that is the route-ahead window paying out, N-1 layers of compute before the +// committing topk even runs. Only the latest submission is accepted; a ranking the previous token +// never collected fails the seq check and dies here rather than routing anything. +void RouterHook::route_ahead_collect(int il) { + if (route_ahead_ <= 0 || ra_tripped_ || batch_phase_ != 1 || !source_) return; + PredictResult r; + { + // Never waited for, only popped — measured, not a stylistic choice: a bounded cv wait here + // (to force every committed layer's reads to start at the predicting layer) bought I/O + // 0.060→0.040 s/token on the host and paid +0.045 s/token of eval-thread wall for it, + // because a cv wakeup costs on the order of a millisecond, i.e. the demand read it was + // trying to save. The depth knob is the honest lever instead: at N>=2 the ranking is + // simply ready a whole layer before its topk, and the next callback issues it with no one + // waiting on anyone. + std::lock_guard lk(pred_mtx_); + if (!pred_result_.ready || !pred_result_.commit || pred_result_.seq != pred_seq_) return; + r = std::move(pred_result_); + pred_result_.ready = false; + } + if (r.nl <= 0 || r.nl >= n_layer_ || r.ids.empty()) return; + ra_pred_[r.nl] = std::move(r.ids); + // r.spec is the drop-aware issue list the worker built alongside the ranking: the committed + // experts the drop policy would actually read. Issue only those; the rest stay cold and die + // at the commit's drop exactly as the router's own cold tail does. + if (r.nl > il) route_ahead_issue_for(r.nl, r.spec); +} + +// At the topk of layer il, before anything reads the ids: replace the router's selection with the +// ranking made N layers back. This runs BEFORE the weight chain's get_rows consumes the ids, so +// the graph itself computes the substituted experts' weights from this layer's true logits — the +// committed routing carries real, renormalized probabilities, not copies of the prediction's. +// Rewriting here also means the gather below, the trace, the drop policy and load_layer all see +// the committed ids: the whole pipeline treats them as THE routing, which is the experiment. +void RouterHook::apply_route_ahead(ggml_tensor * ids, int il, int nu, int nt) { + if (ra_tripped_ || il < 0 || il >= n_layer_ || nu <= 0 || nu > predict_max_k) return; + if (il < route_ahead_) return; // structurally no prediction reaches these layers; not a miss + std::vector & pred = ra_pred_[il]; + // One-token decode rows only: the prediction was made from a single hidden-state row. A wider + // decode batch has no per-token prediction to substitute, so it keeps the router's choice. + if (nt != 1 || (int) pred.size() < nu) { + ++ra_passthrough_; + pred.clear(); + return; + } + // The agreement between the committed selection and the router's own choice IS the measured + // perturbation this flag trades on; score it before the overwrite destroys the evidence. + int hits = 0; + for (int k = 0; k < nu; ++k) { + const int32_t actual = *id_at(ids, 0, k); + for (int p = 0; p < nu; ++p) + if (pred[(size_t) p] == actual) { + ++hits; + break; + } + } + ra_slots_ += nu; + ra_hits_ += hits; + ++ra_overridden_; + for (int k = 0; k < nu; ++k) + *id_at(ids, 0, k) = pred[(size_t) k]; + pred.clear(); +} + +// Start reading layer nl's experts ahead of its topk. `cand` is the worker's drop-aware issue +// list: committed experts the drop policy would actually read, confidence-ordered. Within that +// list no spec_max cap applies — these are not guesses, and a miss not issued here is the same +// read paid on demand later, minus the window. Residents are retained instead (protecting a +// slice that is certainly about to be routed costs zero bytes). Runs on the eval thread, from a +// collect point that sits after some layer's own load was handed to the source — the same +// post-load issue window every other speculation uses, so the demand reads in flight keep their +// head-of-line position. +void RouterHook::route_ahead_issue_for(int nl, const std::vector & cand) { + if (route_ahead_ <= 0 || batch_phase_ != 1 || !source_) return; + if (nl <= 0 || nl >= n_layer_ || cand.empty()) return; + const auto t0 = std::chrono::steady_clock::now(); + ra_ids_.assign(cand.begin(), cand.end()); + // Settle landed speculation first, or a slice already read for an earlier token would classify + // as a miss and be read again — the exact double-spend settle_spec exists to prevent. + source_->settle_spec(); + ra_res_.assign(ra_ids_.size(), (uint8_t) 0); + source_->query_residency(nl, ra_ids_.data(), (int) ra_ids_.size(), ra_res_.data()); + ra_keep_.clear(); + ra_spec_.clear(); + for (size_t i = 0; i < ra_ids_.size(); ++i) + (ra_res_[i] != route_miss ? ra_keep_ : ra_spec_).push_back(ra_ids_[i]); + if (!ra_keep_.empty()) source_->retain(nl, ra_keep_.data(), (int) ra_keep_.size()); + if (!ra_spec_.empty()) source_->prefetch(nl, ra_spec_.data(), (int) ra_spec_.size()); + ra_issue_ns_ += + (long long) std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); +} + // Score one prediction against one routing. Only the prediction's first k entries count: a // prefetch would fetch k experts, so crediting a hit found deeper in the ranking would measure a // predictor nobody could build. @@ -882,8 +1268,9 @@ bool RouterHook::on_eval(ggml_tensor * t, bool ask) { // for free, so it asks for nothing and validates its barrier-less read with the watchdog // instead. Decode-only either way. The "-" in the pattern is what keeps // "ffn_moe_logits_biased-" — a different node — from matching. - const int gl = - (moe_node && predict_on() && source_ && batch_phase_ == 1) ? match_layer_node(t->name, "ffn_moe_logits-") : -1; + const int gl = (moe_node && (predict_on() || route_ahead_ > 0) && source_ && batch_phase_ == 1) + ? match_layer_node(t->name, "ffn_moe_logits-") + : -1; const bool is_logits = gl >= 0; if (ask && is_logits && gl < n_layer_) { ggml_tensor * w = t->src[0]; @@ -908,13 +1295,18 @@ bool RouterHook::on_eval(ggml_tensor * t, bool ask) { } } } + // Route-ahead deliberately does NOT isolate the gate matmul: it wants the node's source + // pointers (free, from this ask pass) and reads the row barrier-less at the topk, exactly + // as the prefetch does — the isolated variant measured ~+0.04 s/token of pure barrier and + // GEMV tax on the host. What makes that safe for a COMMITTED consumer is the watchdog in + // route_ahead_submit plus the passthrough default, not a barrier. return ctrace_iso || is_topk || weights_iso || (is_logits && predict_log_); } // The probe attaches to the gate matmul rather than to the topk node because this is where the // router's own two inputs are reachable: a graph intermediate cannot be found by name later, and // its pointer does not survive to the next graph. It writes nothing the graph will read. - if (is_logits) predict_at_logits(t, gl); + if (is_logits && predict_on()) predict_at_logits(t, gl); // Weights follow their layer's topk, so the pending record is already open; keep the last // one offered (match_weights explains why) and let the flush read it. This runs BEFORE the drop @@ -941,8 +1333,24 @@ bool RouterHook::on_eval(ggml_tensor * t, bool ask) { // The previous layer's weight chain has been fully offered by now. close_drop_layer(); - gathered_.clear(); const int nu = (int) t->ne[0], nt = (int) t->ne[1]; + // Route-ahead, in submission order: first the watchdog sample and the ranking job for + // layer il+N (both need the ROUTER's ids, still untouched here), then the commit — BEFORE + // the ids are gathered, so everything downstream — the trace, the drop policy, load_layer, + // the weight chain the graph is about to run — sees the committed routing, not the + // router's, and nothing disagrees about which experts this layer used. + if (route_ahead_ > 0 && batch_phase_ == 1) { + // Collect FIRST: the ranking racing this topk is keyed to the seq of the job the + // PREVIOUS topk submitted, and the submit below bumps that seq — collecting after it + // would stale-fail nearly every legitimate result (measured: half the layers fell to + // passthrough exactly that way, surviving only when a slow load let an earlier + // collect site catch them). + route_ahead_collect(il); + route_ahead_submit(t, il, nu, nt); + apply_route_ahead(t, il, nu, nt); + } + + gathered_.clear(); for (int j = 0; j < nt; ++j) for (int k = 0; k < nu; ++k) gathered_.push_back( @@ -992,6 +1400,7 @@ bool RouterHook::on_eval(ggml_tensor * t, bool ask) { // Deferred layers speculate from apply_drop instead — after THEIR load, for the same // reason this call sits after the one above: the load's quiesce would cancel it. predict_after_load(il); + route_ahead_collect(il); } // Score the predictors against the routing the router just produced — before the record diff --git a/core/src/moe/router_hook.h b/core/src/moe/router_hook.h index 0ccb34d..92d4b1d 100644 --- a/core/src/moe/router_hook.h +++ b/core/src/moe/router_hook.h @@ -33,6 +33,7 @@ #include "bmoe/predict_stats.h" #include "expert_stream_source.h" +#include #include #include #include @@ -111,6 +112,46 @@ public: // predicted residents happens at every value — it costs zero bytes. void set_predict_prefetch(bool on, int spec_max = 2); + // ── route-ahead: commit to the stale-gate prediction (lossy; see MoeStreamConfig) ─ + // With n > 0, the expert selection of MoE layer L is REPLACED, during decode, by the ranking + // layer L's own gate matrix produced on the hidden state N layers earlier in the same forward + // pass — the prediction the probe scores, acted on as the routing itself. The router still + // computes: its logits feed the weight chain (so the substituted experts carry their true, + // renormalized routing weights) and the layer's own gate input feeds the prediction for L+N. + // Layers 0..N-1, prefill, and any layer whose prediction is unavailable route normally. + // The point: the experts of layer L are known a full N layers of compute before L runs, so a + // prefetch of them is never late and never wrong — the 100%-hit regime no honest predictor + // reaches, bought with a quality perturbation this flag exists to measure. + void set_route_ahead(int n); + + // How hard route-ahead actually bit, and how far the committed choice sat from the router's + // own: `overridden` routings had their ids rewritten, `passthrough` were eligible (decode, + // layer >= N) but had no prediction to commit — the first decode token, plus any layer whose + // gate matrix or input the hook could not read. hits/slots is the agreement between the + // committed selection and what the router would have chosen, the same slot metric the probe + // reports — here it is the measured routing perturbation, not an accuracy claim. + long long route_ahead_overridden() const { return ra_overridden_; } + long long route_ahead_passthrough() const { return ra_passthrough_; } + long long route_ahead_slots() const { return ra_slots_; } + long long route_ahead_hits() const { return ra_hits_; } + + // CPU nanoseconds the prediction worker spent ranking, and how many rankings it produced. + // This is the feature's own arithmetic — one gate GEMV per MoE layer per token, in this + // file's scalar loop rather than ggml's kernels — and it is the one cost that does NOT + // appear in wall time on a machine with spare cores, which is exactly why it has to be + // reported: on a 4-thread phone the same work competes for cores and DRAM bandwidth with + // the decode it is trying to accelerate. + long long route_ahead_gemv_ns() const { return ra_gemv_ns_.load(); } + long long route_ahead_gemv_jobs() const { return ra_gemv_jobs_.load(); } + // Eval-thread nanoseconds spent ISSUING the early reads (settle, residency query, retain, + // prefetch — page commits included). Separate from the GEMV because they live on different + // threads: the GEMV is off the critical path by construction, this is not, and on a phone the + // two are the whole per-layer tax route-ahead adds to the decode. + long long route_ahead_issue_ns() const { return ra_issue_ns_; } + // Eval-thread nanoseconds the sampled fresh-gate watchdog spent on its exact float GEMV — the + // third and last route-ahead cost on the eval thread, metered like the other two. + long long route_ahead_wd_ns() const { return ra_wd_ns_; } + // ── route trace (diagnostics; see bmoe/route_trace.h) ──────────────────────────── // When on, the hook additionally asks for each layer's router-weight node and records one // RouteTraceRow per routed expert. Rows buffer in RAM — the callback runs on a compute @@ -263,6 +304,7 @@ private: uint64_t seq = 0; int nl = 0; // layer the prediction is FOR int nu = 0; + bool commit = false; // route-ahead: rank for commitment instead of building spec lists float drop_frac = 0.0f; int spec_max = 2; const ggml_tensor * gate = nullptr; // layer nl's gate matrix, snapshotted with the job @@ -272,6 +314,8 @@ private: struct PredictResult { uint64_t seq = 0; int nl = -1; + bool commit = false; // produced by a commit job; only route-ahead may consume it + std::vector ids; // commit jobs: the full ranking the topk of nl will commit std::vector spec; // predicted misses to speculate, confidence-ordered, capped std::vector keep; // predicted residents to retain (LRU-protect) bool ready = false; @@ -297,6 +341,48 @@ private: static constexpr double wd_min_frac = 0.98; // below this, the read is not trustworthy static constexpr long long wd_min_slots = 2048; // do not judge before this many slots + // Route-ahead (inert unless route_ahead_ > 0). ra_pred_[L] is the ranking committed at layer + // L's topk, produced from the gate input of layer L-N in the same forward pass — filled and + // consumed within one token, with a seq check so a result the previous token never collected + // cannot be committed by this one. It rides the prefetch's whole barrier-less machinery: the + // gate matrix and input pointers stashed in the ask pass, the row copied at the topk callback, + // the GEMV on the shared prediction worker (predict_prefetch is mutually exclusive, so the + // single job slot is never contended), and the same sampled fresh-gate watchdog validating the + // read. The one thing a COMMITTED prediction changes is the failure mode: where a corrupt + // prefetch wastes a read, a corrupt commit would BE the routing — so on a watchdog trip, and + // for any layer whose prediction is not ready in time, the layer keeps the router's own + // choice (counted as passthrough), never a stale or suspect one. + int route_ahead_ = 0; + std::vector> ra_pred_; + long long ra_overridden_ = 0, ra_passthrough_ = 0, ra_slots_ = 0, ra_hits_ = 0; + bool ra_tripped_ = false; // watchdog verdict: read not trustworthy + std::atomic ra_gemv_ns_{0}, ra_gemv_jobs_{0}; // the prediction's own CPU cost + long long ra_issue_ns_ = 0; // eval-thread cost of issuing the reads + long long ra_wd_ns_ = 0; // eval-thread cost of the sampled watchdog + + // ── int8 mirror of the gate matrices, for the prediction only ──────────────────── + // The prediction re-reads a whole gate matrix per layer per token (n_embd x n_expert, ~4 MB + // at F32) and only needs to ORDER 256 experts to take the top few. Measured on a phone, that + // read is what the GEMV costs: ~160 MB/token of DRAM traffic, ~17-19 ms — bandwidth, not + // arithmetic, which is why vectorising the loop three times over bought little. An int8 + // mirror with one scale per expert quarters the traffic for a ranking that survives it + // trivially (the router's own logits differ by far more than an int8 step, and the sampled + // fresh-gate watchdog is already the alarm if a model ever proves otherwise). Built lazily, + // once per layer, from the same weight leaf the graph uses; ~1 MB per layer of extra RAM. + // The graph's own routing never touches this — only the prediction does. + std::vector> ra_gate_q_; // per layer: n_expert rows of n_embd int8 + std::vector> ra_gate_s_; // per layer: one dequant scale per expert + std::vector ra_qh_; // scratch: the activation row, quantized per call + bool quantize_gate(int il); // build the mirror for layer il; false if unsupported + bool gate_scores_q(int il, const std::vector & h, std::vector & out); + std::vector ra_ids_, ra_keep_, ra_spec_; // scratch: the committed future, split by residency + std::vector ra_res_; // scratch: residency of the committed ids + void route_ahead_submit(ggml_tensor * ids, int il, int nu, int nt); // watchdog + row copy + job + void apply_route_ahead(ggml_tensor * ids, int il, int nu, int nt); + void route_ahead_collect(int il); // pop a finished ranking; issue its reads if still ahead + // Start reading layer nl's committed selection — `cand` is the worker's drop-aware issue list. + void route_ahead_issue_for(int nl, const std::vector & cand); + bool predict_on() const { return predict_log_ || predict_prefetch_; } void predict_reset(); void predict_worker_stop(); diff --git a/docs/README.md b/docs/README.md index 8fd51ed..63b5b08 100644 --- a/docs/README.md +++ b/docs/README.md @@ -25,6 +25,7 @@ for the idea the project is built on. | [expert-dropping.md](expert-dropping.md) | `--drop-cold-experts F`: spending quality only where it buys a flash read, and why it is the one setting whose output is not reproducible. | | [mtp.md](mtp.md) | `--mtp`: drafting with the model's own MTP head and verifying a whole group per decode — lossless by construction, and why a wider decode can lose in the streamed regime. | | [expert-prediction.md](expert-prediction.md) | `--predict-log`: how much of a routing can be known a layer early, measured against the predictor `--prefetch` already bets on — and why a good score still would not mean a faster decode. | +| [route-ahead.md](route-ahead.md) | `--route-ahead N`: committing the routing to the N-layers-early prediction, so a prefetch of it can never miss — and what that costs in quality (experimental, lossy). | | [android-memory.md](android-memory.md) | What reclaims the engine's memory on a phone, which levers exist (almost none), and why the cache hit rate is what the kernel judges you by. | | [pressure.md](pressure.md) | Cache policy under memory pressure: why an unaffordable budget starts a reclaim war, why the adaptive governor was retired, and what the fixed `--cache-mb` / `--dense-weights` levers do. | diff --git a/docs/route-ahead.md b/docs/route-ahead.md new file mode 100644 index 0000000..2b461ae --- /dev/null +++ b/docs/route-ahead.md @@ -0,0 +1,212 @@ +# Route-ahead: committing the routing to its own prediction + +`--route-ahead N` — experimental, lossy, off by default. + +## The idea + +Every prefetch this engine has tried lives with the same ceiling: the routing of MoE layer L is a +function of the hidden state *entering* layer L, which does not exist until layer L−1 has run. A +predictor working any earlier is structurally approximate — measured here, evaluating a layer's +gate on the hidden state one layer back agrees with the real routing on roughly 85% of slots, and +each further layer of anticipation costs ~3.5–4pp more ([expert-prediction.md](expert-prediction.md)). +So a speculative prefetch is always paying for some misses, and the anticipation window it buys is +never wider than the accuracy it can afford. + +Route-ahead inverts the bet: if the prediction cannot reach the routing, lower the routing to the +prediction. With `--route-ahead N`, the expert selection of decode layer L is *replaced* by the +ranking layer L's own gate matrix produced on the hidden state N layers earlier in the same +forward pass. The committed selection is by definition known N full layers of compute before layer +L needs its experts — a prefetch of them can start that early and can never be wrong, the +100%-hit regime no honest predictor reaches. The literature's trained variant of this is Pre-gated +MoE (ISCA'24), which fine-tunes the gate to choose the next layer's experts; route-ahead is the +training-free version, and the flag exists to measure whether the model survives it. + +## What still runs, and what changes + +The router itself is not skipped. Layer L's gate matmul still computes, for two reasons: + +- its logits feed the graph's own weight chain, so after the ids are rewritten at the topk node + the substituted experts receive their **true, renormalized routing weights** — the graph's + `get_rows` + norm/softmax runs on the committed ids exactly as it would on the router's own; +- its gate input is what produces the prediction for layer L+N — read barrier-less at the topk + callback and ranked on the prediction worker thread, the same machinery `--predict-prefetch` + built (and the same sampled fresh-gate watchdog validating the read). Nothing is isolated and + no graph barrier is added, so the eval thread pays only a row copy per layer. Because a + committed prediction's failure mode is the routing itself rather than a wasted read, two + guards close the gap: a watchdog trip disarms the policy out loud and permanently for the run, + and a layer whose ranking is not ready when its topk fires keeps the router's own choice + (counted as passed through) — never a stale or suspect one. At N=1 the worker's deadline is a + single layer of FFN compute, so the ranking usually lands exactly at the committing topk and + the early-read window is thin; at N=2 and above the ranking is collected a layer or more + before its topk and the reads start there — depth buys window, at the staleness cost the + agreement number reports. + +What changes is only *which* experts the selection names. The rewrite happens at the topk node, +before anything consumes the ids, so the trace, the drop policy, `load_layer` and the expert +matmul all agree on the committed routing — the pipeline has no second opinion to leak through. + +Layers `0..N-1` (nothing precedes them by N), prefill, the first decode token (the gate matrices +ahead are not stashed yet) and any layer whose gate or input the hook cannot read route normally +and are counted as passed through, never silently substituted. + +## What it costs + +This changes the output. The committed selection disagrees with the router on a measured fraction +of slots (the `moe-route-ahead:` line reports it per run; expect ~15% at N=1 on a 128-expert +model, worse as N grows), and unlike [expert dropping](expert-dropping.md) — which only removes +low-weight experts from the tail of a routing — a disagreement here can replace a *top-weighted* +expert, and the perturbed hidden state feeds the next layer's routing, so errors can compound +through the depth of the model. Whether they compound into visible damage is exactly the open +question; it cannot be answered from routing logs, only by comparing generations. + +It is therefore mutually exclusive with `--predict-log` (the probe would grade a predictor against +a routing rewritten from that same predictor), and with both prefetchers (`--prefetch`, +`--predict-prefetch`: betting speculative lanes on a future the policy has already fixed is the +same bytes twice). It composes with the cache and with `--drop-cold-experts`, which operate +downstream of whatever the selection says. + +## Reading a run + +``` +moe-route-ahead: 1 layers early — 6096 routings committed, 48 passed through; +committed selection agreed with the router on 84.7% of slots +``` + +The agreement percentage is the size of the perturbation the generation ran under, not an accuracy +claim — with the routing committed there is no miss to speak of. Judge the flag on the generated +text against `--route-ahead 0` on the same prompt with greedy sampling, and on tok/s only after +quality survives. + +## The early reads + +With the LRU cache on, the committed selection is also *acted on*: the moment layer L's own load +is handed to the source, the committed ids of layer L+N — fixed at L's gate matmul a few +callbacks earlier — are split by residency and issued on the same speculative path `--prefetch` +uses: residents are retained (LRU-protected, zero bytes), misses are read on the idle lanes. No +`predict-spec-max`-style cap applies, deliberately: these are not guesses, they are the routing +the topk of L+N will commit verbatim, so a miss not issued early is the same read paid on demand +later, minus the N-layer window. The `moe-prefetch:` line reports them with a `[route-ahead]` +tag; expect the useful fraction to sit at ~100%, which for every other predictor would be +suspicious and here is the construction. Without the cache the selection still commits but +nothing is read early — the flag then only measures quality. + +The issue list is drop-aware, and on a device this is not a nicety but the difference between a +win and a rout: with `--drop-cold-experts` armed, the baseline never reads the routing's cold +low-weight tail at all, and an early read that makes those experts resident un-drops them — a +perfect prefetch of bytes the run was never going to pay for. Measured on a phone at depth 4, +drop 0.75: 78 → 244 MB/token and half the tok/s. So committed experts whose predicted routing +weight sits below the drop threshold are left cold on purpose (the same softmax filter +`--predict-prefetch` applies), and the commit-time drop discards them unread exactly as it does +the router's own tail. With dropping off the filter is inert and the full committed top-k is +issued. + +Committed reads are never cancelled. The speculation machinery was built for guesses, and its +housekeeping treated every queued read as disposable: each demand load (and each settle before an +issue) quiesced the lanes, discarding in-flight and queued speculation wholesale — measured, that +destroyed every early read at depth 2 (0% useful, the same bytes re-bought on demand). Under +route-ahead the queue holds no guesses, so the same housekeeping adopts instead: a loading layer +moves its own committed jobs to the front and finishes them (they are its demand reads, already +staged), completed entries integrate into the cache, and other layers' committed reads stay +queued for the idle lanes. On the serial path the realized window is still one layer — lane time +only exists during a layer's FFN — so depth beyond 1 buys nothing there; under `--overlap` depth +is the lever: on the host A/B (fixed cache, same prompt) N=1 cut the overlap stall from 0.060 to +0.047 s/token and N=2 cut it to 0.010 (−83%) with 97% of early reads useful, worth +20% tok/s on +a machine whose decode is otherwise compute-bound. N=2 is the recommended overlap setting; +deeper adds routing perturbation faster than it adds window. + +## What it costs, measured + +The feature reports its own overhead, because every attempt to reason about it from wall time +alone was wrong. `moe-route-ahead:` prints the prediction's worker CPU (`ra_gemv_ms/tok` in the +CSV trailer) and, separately, the eval-thread cost of issuing the early reads — different threads, +and only the second is on the critical path. Alongside them the cache reports `evictions` and +`re-reads` (a read of an entry it had held before), which is where any prefetch that raises the +byte count while calling its reads useful must be hiding. + +Two of those numbers were bugs, found by instrumenting rather than by argument: + +- **The issue path took the I/O lanes' mutex once per expert.** `prefetch` filtered residency + under a fresh lock per candidate, roughly 120 acquisitions a token on the mutex four lanes pull + every read index from. A blocked eval thread is a descheduled compute thread: the same run with + one lane, hence no contention, did identical work at a quarter of the compute time (0.189 vs + 0.793 s/token). Filtering under one lock, plus an atomic pre-check in the settle, put the whole + issue path at **0.9 ms/token**. +- **The gate GEMV was bandwidth-bound, not compute-bound.** Ranking 256 experts re-reads the + target layer's whole gate matrix — ~4 MB per layer, ~160 MB/token of DRAM traffic on a device + whose decode already saturates the bus, which is why vectorising the loop changed little. An + int8 mirror of each gate (one scale per expert row) with the activation quantized per call + makes the inner loop int8×int8 in int32 — the shape ARM's dotprod instructions exist for — and + cuts it to **6 ms/token** on device from 19, with the committed selection agreeing with the + router on 74.5% of slots against 74.7% for the exact weights. It is aarch64-only and measured + both ways: on x86 the float path is already vectorized and bandwidth is ample, so the int8 + detour costs more than it saves. + +## Status + +**Host: a measured win.** Quality holds — the committed generations are indistinguishable from the +baseline's on a four-prompt objective battery (arithmetic, translation, code, facts: 4/4 correct +in both, one answer byte-identical), on a 512-token essay, and on a second model of a different +generation and quantization. Output is **deterministic**, byte-identical across repeated runs, +unlike [expert dropping](expert-dropping.md), because the prediction depends only on hidden states +and never on what the cache happened to hold. Throughput over three interleaved pairs under +`--overlap` at a fixed cache: **+17%** median with the overlap stall down 0.060 → 0.010 s/token; +**+22%** on the second model; **+40%** with enough cache headroom that nothing is evicted. + +**Device: +21%, measured.** Byte accounting is settled first — an I/O trace of every physical +read, decode-only and equal length, shows the committed routing reads what the baseline reads +(133.6 vs 134.4 MiB/token) with the same 4% read redundancy and zero speculative-then-demand +pairs, so nothing is fetched twice. On throughput, absolutes on this phone are worthless (after a +long session its own baseline swung between 0.79 and 4.54 tok/s in the same cell), so the figure +comes from **interleaved cells** — baseline and N=2 alternating, two passes, same prompt and +length, so thermal drift hits both equally: + +| cell | tok/s | compute | stall | mgmt (of which adopt) | flash | +| --- | --- | --- | --- | --- | --- | +| baseline, 4 lanes | 4.36 | 0.135 | 0.067 | 0.027 | 630 MiB/s | +| N=2, 4 lanes | **5.29** | 0.146 | 0.010 | 0.035 (0.008) | 637 MiB/s | +| N=2, 2 lanes | 5.16 | 0.127 | 0.014 | 0.053 (0.027) | 1045 MiB/s | +| N=2, 1 lane | 2.47 | 0.139 | 0.026 | 0.149 (0.124) | 1090 MiB/s | + +The win is the stall and nothing else: the compute residual barely moves, so this is not a trade +of one cost for another. Lane count is the interesting column — flash throughput saturates well +before four lanes (630 vs 1045 MiB/s of effective bandwidth for the same bytes), so the extra two +lanes buy no bandwidth while costing DRAM traffic that the compute now competes with; at two lanes +the compute residual is the lowest of any cell, but the adoption wait grows to pay it back, and at +one lane the adoption wait dominates outright. Four and two lanes come out level. + +**Inside the demo app the same build measures only +4%, and the cause is not the engine.** With +the app in foreground the device caps its six main cores at 1.9 GHz (2188800 → 1900800 Hz, +measured with the app stopped and running). That slows exactly the compute this flag works to +keep busy, so the share of a token that was stall — the part route-ahead removes — shrinks. It +penalises the whole engine, not this feature. + +The flag stays **off by default**: it is lossy by construction, and a lossy default is a decision +for a release note rather than a benchmark. + +## Next: the second router matmul, and the fork it argues for + +One structural cost remains, and it is not in the algorithm — it is in what this engine may touch. +The baseline computes **one** gate matmul per layer, `gate_L × h_L`. The algorithm as designed +also wants one, substituted: `gate_{L+N} × h_L`. This implementation runs **two**, because the +graph's own router node computes regardless and cannot be removed through the public eval callback. +That second matmul *is* the 6 ms. + +Removing it means changing the graph, i.e. the [fork](seam.md) that already exists for the +expert-ready hook `--overlap` depends on. The shape: + +1. at layer L the graph computes its logits against layer L+N's gate matrix rather than its own; +2. the MoE node takes its expert ids **and weights** from an override hook — the same shape as the + readiness hook — instead of its own argsort and softmax. + +That is roughly 50-100 lines in `build_moe_ffn` plus the engine side, and it takes the prediction's +cost to zero: route-ahead would then compute exactly what the baseline computes. The trade to +measure is quality, not speed: today the substituted experts carry the true router's renormalized +weights (the graph computes them anyway, so they are free); under the substitution they would carry +the prediction's own, which the quality battery above is already set up to judge. + +Now that the device number exists, the fork can be sized honestly: 6 ms against a token that takes +roughly 190 ms there is **about 3%**. It is real and it is the only structural cost left in the +algorithm, but it is not where this feature's remaining time is. The measured costs that are +larger sit in the streaming path both configurations share — the read lanes spend most of their +CPU inside the kernel's O_DIRECT path, and the cache's page commit/evict churn shows up as +`mgmt_ms` — and neither is specific to route-ahead. Those come first. diff --git a/docs/telemetry.md b/docs/telemetry.md index 51655b5..b0142be 100644 --- a/docs/telemetry.md +++ b/docs/telemetry.md @@ -116,8 +116,24 @@ moe-prefetch: MiB speculative, / experts useful ( `` the experts fully read ahead, and `` how many of those a later routing actually hit. See [prefetch.md](prefetch.md). With `--predict-prefetch` the same line is emitted with a `[stale-gate]` tag — same counters, different predictor (see -[expert-prediction.md](expert-prediction.md)). The CSV preamble records which was active -(`prefetch=` / `predict_prefetch=`). +[expert-prediction.md](expert-prediction.md)) — and with `--route-ahead N` with a `[route-ahead]` +tag, where the useful fraction sits at ~100% by construction: the "speculated" ids are the +committed routing itself (see [route-ahead.md](route-ahead.md)). The CSV preamble records which +was active (`prefetch=` / `predict_prefetch=` / `route_ahead=`). + +With `--route-ahead N` a `moe-route-ahead:` line is added: + +``` +moe-route-ahead: layers early — routings committed, passed through; +committed selection agreed with the router on % of slots +``` + +`` decode routings had their expert selection replaced by the N-layers-early +prediction; `` were eligible but had no prediction to commit (the first decode token, and +any layer the hook could not read) and kept the router's choice. The agreement percentage is the +measured routing perturbation: 100% minus it is the fraction of routed slots that went to an +expert the router did not choose. See [route-ahead.md](route-ahead.md). The CSV preamble records +the flag (`route_ahead=`). With `--mtp` or `--ngram` a summary line is added, prefixed with the source that drafted: @@ -202,7 +218,7 @@ prints just the summary lines. n_ctx= n_ubatch= chatml=<0|1> # moe_stream=<0|1> cache_mb= cache_auto=<0|1> cache_floor_mb= cache_ceil_mb= force_cache=<0|1> load_all=<0|1> io_threads= o_direct=<0|1> overlap=<0|1> io_two_wave=<0|1> prefetch= - predict_prefetch=<0|1> predict_log=<0|1> predict_spec_max= prefetch_sync=<0|1> + route_ahead= predict_prefetch=<0|1> predict_log=<0|1> predict_spec_max= prefetch_sync=<0|1> dense_weights= drop_cold_frac= drop_renorm=<0|1> drop_prefill=<0|1> # temp= top_k= top_p= seed= compute_trace_layers= spec= spec_draft_max= mtp_p_min= ngram_min_match= @@ -247,7 +263,8 @@ next to the `turn` column. ``` step,steps,wall_ms,io_ms,compute_ms,read_bytes,cache_hit_pct,stall_ms,mgmt_ms,majflt,cpu_ms, dense_resident_frac,turn,majflt_mib,cache_budget_mib,rss_mib,rss_anon_mib,rss_file_mib,swap_mib, -mem_available_mib,mem_free_mib,swap_free_mib,loop_overhead_ms,mtp_batch,mtp_draft_ms +mem_available_mib,mem_free_mib,swap_free_mib,loop_overhead_ms,mtp_batch,mtp_draft_ms,drain_ms, +adopt_ms,ra_issue_ms,ra_wd_ms ``` `stall_ms`, `mgmt_ms`, `majflt`, `cpu_ms` and `dense_resident_frac` are trailing columns appended @@ -270,7 +287,12 @@ actually cost — `tokens / mtp_decodes` is the amortisation achieved) and `mtp_ that amortisation cost outside the decode, which `s/tok` and `tok/s` both exclude) and `drafted_steps=` (how many steps drafted at all — below `mtp_decodes` only for `--ngram`, which abstains); see the `io_ms` -note above for how the read-time columns are reinterpreted under overlap. +note above for how the read-time columns are reinterpreted under overlap. The route-ahead keys `ra_committed`, +`ra_passthrough`, `ra_agree_pct`, `ra_gemv_ms/tok`, `ra_issue_ms/tok` and `ra_wd_ms/tok` mirror the +CLI's `moe-route-ahead:` report ([route-ahead.md](route-ahead.md)); `drain_s/tok` / `adopt_s/tok` +average the two wait columns below, and `evictions` / `rereads` are the cache-churn counters from +the `moe-cache:` line — a read of an entry the cache had already held once is the only way a +prefetch whose reads are all "useful" can still raise the byte count. The trailing block is the memory picture, added so a run can be diagnosed from its own file: @@ -288,6 +310,10 @@ The trailing block is the memory picture, added so a run can be diagnosed from i | `loop_overhead_ms` | wall time between the **previous** token's decode and this one's: sampling, detokenization, rendering the answer for a UI, and the sink writes. On the first token, the gap from the end of prefill to the first decode. | | `mtp_batch` | how many tokens the decode that produced this row confirmed. `1` without speculation. A verify decode confirms a whole group and its **entire** cost — `wall_ms`, `io_ms`, `majflt`, `cpu_ms`, `read_bytes`, `loop_overhead_ms` — is charged to the group's FIRST row; the rest carry zeros. Without this column those zeros read as free tokens. Per-token cost of a group is the first row's `wall_ms / mtp_batch`. | | `mtp_draft_ms` | time this group spent drafting and catching the draft context up — everything speculation adds **outside** the decode. A slice of `loop_overhead_ms`, not an addition to it. `0` without speculation, and near-zero under `--ngram`, whose drafting is a scan of the token history rather than a decode. This is the column that makes the price of speculation measurable instead of inferable: `wall_ms` never contained it, and neither does `tok/s`. | +| `drain_ms` | overlap only: eval-thread wall waiting for the **previous** layer's read batch to finish before its job slots are reused, at the top of each async load. Part of `compute_ms` — named so a fat compute residual can be attributed instead of theorised about. `0` when it costs nothing. | +| `adopt_ms` | route-ahead only: the load waiting for its own committed speculative reads to complete before staging (adoption). Part of `mgmt_ms`. | +| `ra_issue_ms` | route-ahead only: eval-thread time issuing the committed selection's early reads (settle, residency split, retain, prefetch — page commits included). Part of `compute_ms`. | +| `ra_wd_ms` | route-ahead only: the sampled fresh-gate watchdog's exact float GEMV on the eval thread. Part of `compute_ms`. | All are `0` where the platform cannot report them (the Windows host build reports device memory but not the per-process split). diff --git a/examples/android/app/src/main/java/io/bigmoeonedge/example/AppSettings.kt b/examples/android/app/src/main/java/io/bigmoeonedge/example/AppSettings.kt index c551bb6..f685840 100644 --- a/examples/android/app/src/main/java/io/bigmoeonedge/example/AppSettings.kt +++ b/examples/android/app/src/main/java/io/bigmoeonedge/example/AppSettings.kt @@ -47,6 +47,13 @@ data class AppSettings( // spends no flash and only protects predicted residents from eviction). Defaults to 0 because // the matched-pair A/B showed read-ahead losing on a saturated flash (docs/expert-prediction.md). val predictSpecMax: Int = 0, + // Route-ahead (experimental, LOSSY): decode routing at layer L is COMMITTED to the prediction + // made N layers earlier in the same forward pass, and with the cache on the committed experts + // are read that early — reads that can never be wasted, since they ARE the routing. Changes + // the output (~20% of slots re-route at N=1; quality held in the first host A/B). Excludes + // both prefetchers, so sessionArgv only emits it when they are off. 0 = off, the default + // until the on-device A/B earns it more. + val routeAhead: Int = 0, // Cache-aware expert dropping, as a PERCENTAGE of the uniform share 1/top-k (0 = off, 100 = the // share itself). Stored as an Int because the settings are integer rungs; the flag takes a // fraction. LOSSY and cache-dependent — it changes the output, and not reproducibly. @@ -144,6 +151,14 @@ data class AppSettings( a += "--predict-prefetch" a += listOf("--predict-spec-max", predictSpecMax.toString()) } + // Route-ahead excludes both prefetchers (the engine refuses the pairs: they would + // speculate a future route-ahead has already fixed) and self-speculation (a verify + // decode is several positions wide, and route-ahead declines to commit on every one of + // them while still paying for the prediction). It works without the cache too — the + // routing still commits — but only reads early when the cache is on. + if (routeAhead > 0 && prefetchLayers == 0 && !predictPrefetch && spec == SPEC_OFF) { + a += listOf("--route-ahead", routeAhead.toString()) + } // Cache-aware dropping needs a live cache to ask about residency — with the cache off // every expert reads as a miss and the engine rejects the combination outright, so the // same cacheOn condition that guards prefetch guards this. The engine takes a fraction @@ -170,8 +185,8 @@ data class AppSettings( */ fun sessionSignature(modelPath: String): String = listOf(modelPath, mmap, cacheMb, cacheCeilMb, ioThreads, threads, nExpertUsed, sessionCtx, oDirect, - overlap, denseWeights, prefetchLayers, predictPrefetch, predictSpecMax, dropColdPct, - spec, mtpDraft, mtpPMinPct) + overlap, denseWeights, prefetchLayers, predictPrefetch, predictSpecMax, routeAhead, + dropColdPct, spec, mtpDraft, mtpPMinPct) .joinToString("|") fun save(ctx: Context) { @@ -186,6 +201,7 @@ data class AppSettings( .putInt("prefetchLayers", prefetchLayers) .putBoolean("predictPrefetch", predictPrefetch) .putInt("predictSpecMax", predictSpecMax) + .putInt("routeAhead", routeAhead) .putInt("dropColdPct", dropColdPct) .putInt("sessionCtx", sessionCtx) .putString("spec", spec).putInt("mtpDraft", mtpDraft).putInt("mtpPMinPct", mtpPMinPct) @@ -287,6 +303,10 @@ data class AppSettings( // default — the matched-pair A/B showed 2 losing −21% on a saturated flash; anything above // it re-buys the measured full-speculation pathology. val PREDICT_SPEC_CHOICES = intArrayOf(0, 1, 2, 4) + // Route-ahead depth: how many layers early the routing is committed (and read). Each layer + // of depth widens the I/O window and the routing perturbation together; 1 is the measured + // sweet spot on the host, 4 the edge where damage first showed. + val ROUTE_AHEAD_CHOICES = intArrayOf(0, 1, 2, 4) // Percent of the uniform share 1/top-k. 100 is the share itself and the useful maximum: // above it the threshold could exceed every weight in a routing. The rungs below it are the // conservative half of the curve, where the replay already beats a top-k cut on both axes. @@ -321,6 +341,7 @@ data class AppSettings( prefetchLayers = p.getInt("prefetchLayers", d.prefetchLayers), predictPrefetch = p.getBoolean("predictPrefetch", d.predictPrefetch), predictSpecMax = p.getInt("predictSpecMax", d.predictSpecMax), + routeAhead = p.getInt("routeAhead", d.routeAhead), dropColdPct = p.getInt("dropColdPct", d.dropColdPct), sessionCtx = p.getInt("sessionCtx", d.sessionCtx), spec = run { diff --git a/examples/android/app/src/main/java/io/bigmoeonedge/example/MetricFields.kt b/examples/android/app/src/main/java/io/bigmoeonedge/example/MetricFields.kt index 8a635a6..bdd324e 100644 --- a/examples/android/app/src/main/java/io/bigmoeonedge/example/MetricFields.kt +++ b/examples/android/app/src/main/java/io/bigmoeonedge/example/MetricFields.kt @@ -44,6 +44,11 @@ object MetricFields { MetricField("loop_overhead_ms", "time between tokens", "everything outside llama_decode: sampling, detokenization, rendering the answer, writing the sinks. It falls outside wall_ms and so outside the reported tok/s, which is why it needs a column — work that moves only this number is paid on every token and shows up nowhere else. On the first token, the gap from the end of prefill to the first decode", Better.LOWER), + MetricField("drain_ms", "wait on previous batch", "overlap only: eval-thread time waiting for the PREVIOUS layer's reads to finish before reusing their slots — part of compute_ms, named so a large compute residual can be attributed instead of guessed at", Better.LOWER), + MetricField("adopt_ms", "wait adopting committed reads", "route-ahead only: the load waiting for its own committed speculative reads to complete before staging — part of mgmt_ms", Better.LOWER), + MetricField("ra_issue_ms", "issuing early reads", "route-ahead only: eval-thread time issuing the committed selection's reads (settle, residency, retain, prefetch) — part of compute_ms", Better.LOWER), + MetricField("ra_wd_ms", "watchdog control", "route-ahead only: the sampled fresh-gate control's exact GEMV on the eval thread — part of compute_ms", Better.LOWER), + MetricField("mem_available_mib", "device 'available' (it lies)", "what the kernel claims is free — it counts our own mmap'd weights as reclaimable, so it over-states headroom", Better.NEUTRAL), MetricField("mem_free_mib", "device free", "truly free RAM, before any reclaim", Better.NEUTRAL), MetricField("swap_free_mib", "swap free", "zram space left before the device is truly out of room", Better.NEUTRAL), @@ -89,6 +94,7 @@ object ConfigFields { ConfigField("io_two_wave", "the first projection's reads were published to the lanes as soon as they were staged, rather than after the whole layer"), ConfigField("load_all", "every expert was loaded each token, routing ignored. An A/B baseline"), ConfigField("prefetch", "temporal prefetch depth in layers, betting this token routes like the last. 0 = off"), + ConfigField("route_ahead", "each layer's expert choice was COMMITTED to the prediction made this many layers earlier in the same forward pass, and read that early. Lossy: a fraction of choices differ from the router's own. 0 = off"), ConfigField("predict_prefetch", "the next layer's router ran early, and the cache acted on that prediction instead of the previous token's routing"), ConfigField("predict_spec_max", "how many predicted, non-resident experts were read ahead. 0 = retention only, which spends no flash and merely protects predicted residents from eviction. Recorded but unused when predict_prefetch is off"), ConfigField("predict_log", "prediction-accuracy probe. A diagnostic, not a shipping setting"), diff --git a/examples/android/app/src/main/java/io/bigmoeonedge/example/MetricsCsv.kt b/examples/android/app/src/main/java/io/bigmoeonedge/example/MetricsCsv.kt index b6a138a..4aaa495 100644 --- a/examples/android/app/src/main/java/io/bigmoeonedge/example/MetricsCsv.kt +++ b/examples/android/app/src/main/java/io/bigmoeonedge/example/MetricsCsv.kt @@ -46,6 +46,7 @@ internal class Csv( // dropped experts is not the same KIND of run as one that did not, and two compare // legends that differ only by this used to read as identical (#136). if (info["predict_prefetch"] == "1") "predict" else null, + info["route_ahead"]?.takeIf { it != "0" }?.let { "route-ahead $it" }, info["drop_cold_frac"]?.takeIf { (it.toFloatOrNull() ?: 0f) > 0f }?.let { "drop $it" }, // Same argument, and stronger: under speculation a decode confirms a whole group, so // the per-token rows are not even accounted the same way (mtp_batch). Two runs that diff --git a/examples/android/app/src/main/java/io/bigmoeonedge/example/MetricsScreen.kt b/examples/android/app/src/main/java/io/bigmoeonedge/example/MetricsScreen.kt index 21a6705..6afa109 100644 --- a/examples/android/app/src/main/java/io/bigmoeonedge/example/MetricsScreen.kt +++ b/examples/android/app/src/main/java/io/bigmoeonedge/example/MetricsScreen.kt @@ -410,6 +410,7 @@ private val CONFIG_ORDER = listOf( "io_two_wave" to "Two-wave publish", "load_all" to "Load all experts", "prefetch" to "Temporal prefetch", + "route_ahead" to "Route-ahead (layers)", "predict_prefetch" to "Predictive prefetch", "predict_spec_max" to "Speculated misses/layer", "predict_log" to "Prediction probe", @@ -461,7 +462,7 @@ private fun prettyConfigValue(key: String, v: String, info: Map) // dropping happened" is exactly the misreading this display exists to prevent. The drop // fraction is shown as the engine took it (a fraction of the uniform share 1/top-k, not of the // routing) so it matches --drop-cold-experts and the settings screen. - (key == "prefetch" || key == "drop_cold_frac") && v.toFloatOrNull() == 0f -> "off" + (key == "prefetch" || key == "route_ahead" || key == "drop_cold_frac") && v.toFloatOrNull() == 0f -> "off" key == "predict_spec_max" && v == "0" -> "0 (retention only)" else -> v } diff --git a/examples/android/app/src/main/java/io/bigmoeonedge/example/SettingsScreen.kt b/examples/android/app/src/main/java/io/bigmoeonedge/example/SettingsScreen.kt index 938eabf..0729034 100644 --- a/examples/android/app/src/main/java/io/bigmoeonedge/example/SettingsScreen.kt +++ b/examples/android/app/src/main/java/io/bigmoeonedge/example/SettingsScreen.kt @@ -120,7 +120,7 @@ fun SettingsScreen(current: AppSettings, onChange: (AppSettings) -> Unit, onBack format = { if (it == 0) "off" else "$it" }, // Mutually exclusive with predictive prefetch: two predictors would speculate // the same future twice, and the engine refuses the pair. - enabled = stream && cacheOn && !current.predictPrefetch, + enabled = stream && cacheOn && !current.predictPrefetch && current.routeAhead == 0, ) { onChange(current.copy(prefetchLayers = it)) } Text( "Experimental. Bets a layer will reuse the experts it picked for the previous token, and " + @@ -129,10 +129,12 @@ fun SettingsScreen(current: AppSettings, onChange: (AppSettings) -> Unit, onBack ) SwitchRow( "Predictive prefetch (experimental)", - "Experimental. Rather than betting on the previous token, runs the next layer's own " + - "router early to ask which experts it will actually want. Reads ahead within the " + - "budget below and protects what it names. Needs the cache; replaces the above.", - current.predictPrefetch, enabled = stream && cacheOn && current.prefetchLayers == 0, + "Instead of betting on the previous token, asks the next layer's own router one " + + "layer early to find out which experts it will actually want. Reads ahead within " + + "the budget below and keeps what the prediction names. Needs the cache; replaces " + + "temporal prefetch.", + current.predictPrefetch, + enabled = stream && cacheOn && current.prefetchLayers == 0 && current.routeAhead == 0, ) { onChange(current.copy(predictPrefetch = it)) } if (current.predictPrefetch) { IntSetting( @@ -215,6 +217,22 @@ fun SettingsScreen(current: AppSettings, onChange: (AppSettings) -> Unit, onBack "the reading, and changes the reply — a deliberate trade of quality for speed.", fontSize = 12.sp, color = MaterialTheme.colorScheme.onSurfaceVariant, ) + IntSetting( + "Route-ahead (layers)", AppSettings.ROUTE_AHEAD_CHOICES, current.routeAhead, + format = { if (it == 0) "off" else "$it" }, + // Excludes both prefetchers (the engine refuses speculating a future this has + // already fixed) and guessing ahead, where a verify pass is several tokens wide + // and this declines to commit on all of them while still paying for itself. + // Needs streaming; the cache is what turns it into early reads. + enabled = !current.mmap && current.prefetchLayers == 0 && !current.predictPrefetch && + current.spec == AppSettings.SPEC_OFF, + ) { onChange(current.copy(routeAhead = it)) } + Text( + "Experimental, changes the output. Each layer's expert choice is committed N layers " + + "early — so with the cache on their reads start that early and are never wasted. " + + "~20% of choices differ from the router's at 1 layer; quality held in the host A/B.", + fontSize = 12.sp, color = MaterialTheme.colorScheme.onSurfaceVariant, + ) IntSetting( "Drop cold experts (% of even share)", AppSettings.DROP_COLD_CHOICES, current.dropColdPct, // The rung labels carry the trade, so the blurb below does not have to repeat diff --git a/tests/config_test.cpp b/tests/config_test.cpp index 4bf7d42..61778e7 100644 --- a/tests/config_test.cpp +++ b/tests/config_test.cpp @@ -112,6 +112,53 @@ int main() { expect_ok("prefetch allowed with the cache on", c); } + // Route-ahead: needs streaming, bounded, excludes the probe and both prefetchers. It does NOT + // need the cache — committing the routing is orthogonal to how the committed experts get read. + { + RunConfig c = ok_base(); + c.moe.route_ahead = 1; // enabled stays false + expect_fail("route_ahead requires streaming", c); + } + { + RunConfig c = ok_moe(); + c.moe.route_ahead = 1; + expect_ok("route_ahead with cache off", c); + c.moe.route_ahead = MoeStreamConfig::route_ahead_max + 1; + expect_fail("route_ahead out of range", c); + c.moe.route_ahead = -1; + expect_fail("route_ahead negative", c); + } + { + RunConfig c = ok_moe(); + c.moe.route_ahead = 1; + c.moe.predict_log = true; + expect_fail("route_ahead excludes predict_log", c); + } + { + RunConfig c = ok_moe(); + c.moe.route_ahead = 1; + c.moe.cache_mb = MoeStreamConfig::cache_min_mb; + c.moe.predict_prefetch = true; + expect_fail("route_ahead excludes predict_prefetch", c); + c.moe.predict_prefetch = false; + c.moe.prefetch_layers = 2; + expect_fail("route_ahead excludes prefetch_layers", c); + c.moe.prefetch_layers = 0; + expect_ok("route_ahead with the cache and no predictor", c); + } + // A verify decode is several positions wide, and route-ahead declines to commit on every one of + // them while still paying for the prediction and the early reads — measured, see config.cpp. + { + RunConfig c = ok_moe(); + c.moe.route_ahead = 1; + c.spec.source = DraftSource::mtp; + expect_fail("route_ahead excludes the MTP draft source", c); + c.spec.source = DraftSource::ngram; + expect_fail("route_ahead excludes the n-gram draft source", c); + c.spec.source = DraftSource::none; + expect_ok("route_ahead with speculation off", c); + } + // Sampling ranges — enforced only when temp > 0. { RunConfig c = ok_base();