From b27c69266ba414d28393825d5aaae35354f81c7a Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Sat, 24 Oct 2026 05:25:09 -0300 Subject: [PATCH] serve: --moe-predict off|router|block, and the block predictor in the docs The prefetch that hides a streamed model's expert reads was predicting with the gate-ahead router, which does not fit Qwen3.8-Flash-Next's hyper-connections: layer l+1's router input is its own mixer applied to the residual after layer l's MoE, so the residual-only form mixes through the wrong weights and is a net loss on a drive at its limit (docs/moe-streaming.md, "Flash-Next, streamed"). docs/moe-streaming.md now carries the predictor that fits -- l+1's own token-mixer block on the current residual, before l's MoE output exists, with the recurrent/conv state writes redirected to scratch -- with the flag, its measurement variables and the explicit note that it is untested on device. serve gains `--moe-predict off|router|block`, passes it through to the llama.cpp pin as ONEBIT_MOE_PREDICT beside ONEBIT_MOE_PREFETCH, and refuses a value it does not know. The graph side lives in the pin (src/onebit-moe-predict.*, one hook in the qwen4exp graph), and nothing here changes a running configuration: the option does nothing without --moe-slots, which this build refuses until the HRX streamer lands. Co-Authored-By: Claude Fable 5.1 --- app/serve.cpp | 14 ++++++++++++++ docs/moe-streaming.md | 39 +++++++++++++++++++++++++++++++++++++-- docs/serve.md | 11 ++++++++++- 3 files changed, 61 insertions(+), 3 deletions(-) diff --git a/app/serve.cpp b/app/serve.cpp index db85a2a4..517a0591 100644 --- a/app/serve.cpp +++ b/app/serve.cpp @@ -161,6 +161,7 @@ struct Options { std::string dflash; // DFlash block-diffusion draft model for the target int moe_slots = 0; // --moe-slots N: stream routed experts from the model file, N held in RAM (-1: auto) int moe_prefetch = 0; // --moe-prefetch N: gate-ahead prefetch N MoE layers ahead (0: off) + std::string moe_predict; // --moe-predict off|router|block: what the prefetch predicts with (docs/moe-streaming.md) std::string moe_subst; // --moe-subst R: a resident expert scoring at least R x a missing one takes its place std::string mmproj; // a vision (or audio) projector: image parts in chat messages int mtp_max = 0; @@ -976,6 +977,10 @@ Launch launch_for(const Options& o, const std::string& device, int child_port) { env.push_back("ONEBIT_MOE_SLOTS=" + std::to_string(slots)); // off unless asked: where the drive is the limit, wrong guesses cost more than right ones save env.push_back("ONEBIT_MOE_PREFETCH=" + std::to_string(o.moe_prefetch)); + // what the prefetch predicts with: the gate-ahead router (the streamer's old form, + // wrong on hyper-connections) or the next layer's own block (docs/moe-streaming.md, + // "Predicting the next layer's experts"). Untested on device: off by default. + if (!o.moe_predict.empty()) env.push_back("ONEBIT_MOE_PREDICT=" + o.moe_predict); if (!o.moe_subst.empty()) env.push_back("ONEBIT_MOE_SUBST=" + o.moe_subst); // the experts stay file-backed: without these, llama.cpp copies every one into RAM // (a pinned host buffer, or a CPU repack), which is what streaming avoids @@ -2593,6 +2598,9 @@ void usage(FILE* out) { " N held in RAM (auto: what fits); R: resident experts stand in for missing\n" " ones; prefetch N layers ahead (default 0). Not in this build: it moves\n" " to HRX with Qwen3.8-Flash-Next (docs/moe-streaming.md)\n" + " [--moe-predict off|router|block] what the prefetch predicts the next layer's experts with:\n" + " router (cheap) or block (that layer's own mixer block; untested on\n" + " device). Only with --moe-slots. Default off\n" " [--mmproj MMPROJ.gguf] images in chat messages (hrx, cpu)\n" " [--parallel N] N requests decoded together (continuous batching)\n" @@ -2681,6 +2689,12 @@ int run_serve(int argc, char** argv) { o.moe_slots = v == "auto" ? -1 : std::stoi(v); } else if (a == "--moe-prefetch") o.moe_prefetch = std::stoi(next()); + else if (a == "--moe-predict") { + o.moe_predict = next(); + if (o.moe_predict != "off" && o.moe_predict != "router" && o.moe_predict != "block") { + throw std::runtime_error("--moe-predict takes off, router or block, not '" + o.moe_predict + "'"); + } + } else if (a == "--moe-subst") o.moe_subst = next(); else if (a == "--mmproj") o.mmproj = next(); else if (a == "--mtp-max") o.mtp_max = std::stoi(next()); diff --git a/docs/moe-streaming.md b/docs/moe-streaming.md index 23ff6078..7301bb65 100644 --- a/docs/moe-streaming.md +++ b/docs/moe-streaming.md @@ -194,6 +194,39 @@ Predicting 2 x top-k covers 93-98% of misses, but reads 5-10 times as many exper never used. Predictors that need only the trace (the previous token's experts, co-occurrence with the previous layer) cover under 30%. +### The next layer's own block (Flash-Next; untested on device) + +The gate-ahead predictor above does not fit Qwen3.8-Flash-Next: its layers keep four 2560-wide +residual streams (hyper-connections), so layer l+1's router input is l+1's own mixer applied to +the residual *after* layer l's MoE, and applying l+1's router to layer l's FFN input mixes +through the wrong weights. The measured result is in "Flash-Next, streamed" below: right about +two times in three, and a loss on a drive that is already the limit. + +The predictor that does fit runs layer l+1's own token-mixer block on the current residual, +before layer l's MoE output exists: l+1's PLE injection, its attention or DeltaNet block, its +combine, then its FFN mixer and its router. Only the ranking of the router logits is used, and +more candidates than the layer routes to are kept so the near misses are read too. The output of +the model does not change: the block only decides which experts a prefetch reads, and whatever a +token really routes to is still waited for. + +State writes of that speculative pass are suppressed where the exact pass would read them: the +DeltaNet recurrent state and the conv histories are written into scratch tensors, and an +attention successor needs nothing, because the exact block rewrites the same cache cell before +anything reads it. The speculative nodes are also left out of the fused-op manifest, so they +cannot make a fused dispatch claim the wrong tensor. + +`--moe-predict off|router|block` (env `ONEBIT_MOE_PREDICT`) selects it, and takes effect with +`--moe-prefetch 1` once the HRX llama.cpp runs the streamer; `router` keeps the cheap +residual-only form for comparison. `ONEBIT_MOE_PREDICT_DEPTH`, `_MARGIN` (read a deeper candidate +only while its logit is within the given distance of the best one), `_MIN`, `_LAYERS` and +`_OUT` (a trace of predictions and the routing that was really used) tune and measure it, and +`ONEBIT_MOE_PREDICT_STATS=1` prints the hit rate per layer when the process exits. + +The graph side is `src/onebit-moe-predict.*` plus one hook in the `qwen4exp` graph in the +llama.cpp pin, and it is **untested on device**: the predictor, its readback and the hit-rate +table are built and unit-tested host-side (`tests/test-onebit-moe-predict.cpp`), and the +prefetch consumer (the streamer's in-graph op) is the step that needs the GPU lane. + ## The expert cache in the engine `moe/` holds the cache (Linux only, built into `1bit` as `ONEBIT_MOE`): @@ -633,7 +666,9 @@ the change is within the noise (four rounds, 6.7 against 7.0 tok/s). 2. **A better stand-in than the router's next choice.** On Flash-Next, r 0.5 costs a KL divergence of 0.07 ("Flash-Next, streamed"). Candidates: weighing a stand-in by how often it co-occurs with the missing expert, or reading only the missing expert's most-weighted rows. -3. **Flash-Next's gate-ahead prediction.** Its layers keep four 2560-wide residual streams - (hyper-connections), so the predictor needs that model's mixing step before the router. +3. **Flash-Next's prediction.** Its layers keep four 2560-wide residual streams + (hyper-connections), so the predictor needs that model's mixing step before the router. Built + as the next layer's own block ("The next layer's own block", above); it needs a device run to + rate it and the streamer to consume it. 4. **The quants not on the box** (UD-Q2/Q3/Q4_K_XL): their decode speed, once there is room on the drive. diff --git a/docs/serve.md b/docs/serve.md index 4a03e204..c9279af7 100644 --- a/docs/serve.md +++ b/docs/serve.md @@ -32,7 +32,8 @@ so it also works on its own with any OpenAI client. [--mtp HEAD.gguf | --dflash DRAFT.gguf | --no-mtp] [--mtp-max N] [--mtp-p-min P] [--mmproj MMPROJ.gguf] [--laya | --laya-model DIR] [--route-policy FILE] [--pm [--pm-experts FILE] [--lemonade-url URL]] - [--moe-slots N|auto] [--moe-subst R] [--moe-prefetch N] # refused in this build, below + [--moe-slots N|auto] [--moe-subst R] [--moe-prefetch N] [--moe-predict off|router|block] + # refused in this build, below [--parallel N] [--embed MODEL.gguf] [--rerank MODEL.gguf] [--npu-opt KEY=VALUE ...] @@ -713,6 +714,14 @@ Prompts are slow when streaming: a batch larger than 8 tokens runs the experts o from the mapped file (Flash-Next: 0.4-0.8 tok/s of prompt). Keep prompts short, or see docs/moe-streaming.md, Next. +`--moe-predict off|router|block` (default off; only with `--moe-slots`, and a prefetch to predict +for means `--moe-prefetch 1`) chooses what the prefetch predicts the next layer's experts with. `block` runs the next layer's own mixer block +on the current residual, which is what makes the prediction fit Flash-Next's hyper-connections; +`router` keeps the cheap residual-only form. The predictor changes only which experts are read +early, never the output (`docs/moe-streaming.md`, "The next layer's own block"). **Untested on +device**: the predictor and its hit-rate table are built and unit-tested host-side, and it needs +the HRX streamer to consume it. + ## Many requests at once (`--parallel`) `--parallel N` gives llama-server N slots (`-np N`): requests that arrive together decode