serve: --moe-predict off|router|block, and the block predictor in the docs

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;
@@ -993,6 +994,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
@@ -2618,6 +2623,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"
 
@@ -2706,6 +2714,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());

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.

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 ...]
@@ -768,6 +769,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