Repository navigation
qwen4exp: predict layer l+1's routing by running its own block (ONEBIT_MOE_PREDICT) - #137
bong-water-water-bong wants to merge 1 commit into
Conversation
…T_MOE_PREDICT)
The streamed-expert prefetch on Qwen3.8-Flash-Next paid for a prediction that cannot work on
that model: with hyper-connections, layer l+1's router input is l+1's own mixer applied to the
residual after layer l's MoE, so applying l+1's router to layer l's FFN input mixes through the
wrong weights (docs/moe-streaming.md: right about two times in three, and a net loss on a drive
that is already the limit).
This adds the predictor that does fit. Inside layer l's graph, before l's MoE, run l+1's own
token-mixer block on the current residual -- its PLE injection, its attention or DeltaNet block,
its combine -- then l+1's FFN mixer and its router, and keep the top-k ranking of the logits,
deeper than the layer's own top-k so the near misses are read too. Only the ranking is used, and
whatever a token really routes to is still waited for, so the model's output is unchanged.
State writes of the speculative pass are suppressed where the exact pass would read them: the
DeltaNet recurrent state and the conv histories (DeltaNet and PLE) are written into scratch
tensors instead of the cache. An attention successor needs no suppression: 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 a fused dispatch cannot claim the wrong tensor.
New files, ours: src/onebit-moe-predict.{h,cpp} (environment, the margin gate, the hit-rate table,
the trace, the prefetch sink) and tests/test-onebit-moe-predict.cpp, which checks the policy
host-side. Small hooks: the layer loop in qwen4exp, the state tensors of
build_layer_attn_linear/build_ple, one line in build_moe_ffn to capture the routing that each
layer really used, and one line in llama_context::graph_compute for the readback.
Off by default: ONEBIT_MOE_PREDICT=off|router|block, optionally _DEPTH, _MARGIN, _MIN, _LAYERS,
_NO_ATTN, _OUT (trace), _STATS. Untested on device -- no GPU run was made -- and the prefetch
consumer (the streamer's in-graph op) is the next step; with nothing installed the sink is the
trace, which is the hit-rate measurement the device plan asks for.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
Review (orchestrator): predictor in our own files with a unit test, hooks in the qwen4exp layer loop and the delta-net base kept to parameters and a spec flag, everything behind ONEBIT_MOE_PREDICT (default off). Draft until the device measurement: hit rate vs the trace (router vs block), per-layer cost of the speculative block, tok/s streamed and resident on Flash-Next, after the leaderboard releases the lane. |
|
Triage: needs work — three things, none of them the predictor itself.
Diff read: Related: engine ggml-org#551 already documents the graph side as living "in the llama.cpp pin", which is not true of |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: APPROVE — off by default, the "what is not done" list is honest, and the design record exists where a reader will look. Three things I would want before anyone enables it, the first of which I think is the whole review.
1. The riskiest claim has no gate: "the model's output is unchanged". Everything else in the PR rests on it — a speculative attention/DeltaNet block is executed inside layer l's graph, ahead of the data the exact pass uses, and the only thing standing between that and corrupted output is state-write suppression. That is a claim a test can settle exactly, and cheaply, once a device run happens:
ONEBIT_MOE_PREDICT=off vs ONEBIT_MOE_PREDICT=block -> logits must be BIT-IDENTICAL
Bit-identical, not close, because the PR's own argument is that only the ranking is used: nothing about the arithmetic of the real path changes, so any difference at all is leaked state — a scratch tensors' lifetime, a conv history, a cache cell the exact block did not rewrite. If that assertion is in the device plan, say so and this paragraph is done; if it is not, it is the first row I would add, ahead of hit rate and ahead of pass time. It is also the only one of the three that can be a bug; the other two are tuning.
2. "one scratch per shape serves every layer" — does the scheduler guarantee the order that assumption needs? With hyper-connections there are many predicted layers in one step, and if every speculative pass writes the same shape scratch, then two passes writing it are a write-after-write. In a single graph those nodes are independent, so the order they execute in is the scheduler's choice, not the layer loop's — and if pass A's router reads while pass B has overwritten the scratch, A's ranking is stale. Note the failure mode is the good one: the model output would still be correct (only the prefetch choice is wrong), so it would show up as a lower hit rate, not a wrong logit — which means the measurement you are about to trust could be measuring this instead of the mechanism. If the sharing is safe because the passes are serialised by a dependency I have missed, one sentence in the code comment naming that dependency makes it reviewable; if it is safe only because nothing has reordered yet, a scratch per layer removes the doubt for a few KB.
3. Memory: a speculative block is a whole block. Each predicted layer adds its neighbour's attention or DeltaNet block plus mixer and router to the graph, so its tensors have to be planned and reserved even though their result is discarded. On this box a memory surprise is not a slowdown: it is a load-time refusal or an allocation failure (my own runs today: refusals at 13.6 / 13.5 / 10.8 GiB of pool budget for a 4.7 GiB model, and a 900 s silent stall at 16.7 GiB). So the PR should state what _LAYERS costs in reserved bytes per predicted layer — with _LAYERS subsetting as the mitigation a user is told about. "Predict everything" and "predict eight layers" differ by an order of magnitude here, and the flag makes that the user's decision.
Smaller points, none blocking:
- The
-mddouble-count is disclosed and correctly scoped ("the hit ratio is unaffected" is right: a step counted twice in both numerator and denominator does not move a ratio). Good. register_fused=false, spec=trueis the right instinct — a fused dispatch claiming a speculative node would silently make the trace lie. Worth one clause saying what a non-fused backend does instead: the node is simply materialised, which for a 40 B MoE at 4,608 slots is the arithmetic in point 3.ONEBIT_MOE_PREFETCH's failure is the best part of the write-up: "right about two times in three, and a net loss on a drive that is already the limit — 133 s without prefetch, 247 s with". That is a 1.86× regression measured, and it is the argument for spending a whole block to get the routing right instead of guessing it. If that comparison can be repeated on the same build withroutervsblockvsoff, the PR becomes its own justification.
I did not build or run this (no device arm exists yet, and the model is a 40 B MoE); I am reviewing the mechanism, the suppression argument, and the cost model. Refs the streamed-expert plan in engine docs/moe-streaming.md.
|
VERDICT: CHANGES at Blocking1. Base is So merging this lands on a line with no HRX backend and none of the fleet's work, and it cannot be merged forward either — the histories are unrelated. Retarget to 2. The 3. The acceptance is unmeasured, and the consumer it exists for is not built. The body says so plainly, which is the right way to file it — but the title and the Verified, so the reviewer does not repeat it
Out of scope here, for whoever picks this upThe |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: CHANGES at 5f0d4619ae (posted as a review so head-tracking sees it — the full reasoning is in my comment 6727663931 on this PR).
Blocking, measured on the fork today:
- Base is
master, which the engine does not build.merge-base(1bit/hrx-vulkan-patched, this branch)is empty andmerge-treereports "refusing to merge unrelated histories";master(c35b667) carries nosrc/llama-kv-sharewhile the pinned line does, the two are 10,794 commits apart, and every recent fork PR targets1bit/hrx-vulkan-patched(#233/#235 merged, #240 open). Retarget/rebase there. - The
ctest -R onebitclaim is from 2026-10-24 onmaster. Re-run the host build + that ctest after the rebase, in a graph layer that has moved 10.8k commits (including the fused-node registration this PR edits). - The perf framing is unmeasured and the consumer is unbuilt (the body says so): either add the hit-rate table the flag's defaults are tuned against, or rescope the title/body to plumbing only, off by default, no perf claim.
Verified good, so nobody repeats it: bool register_fused = true is a default in models.h and all six other callers (bailingmoe3, glm5-next, kimi-k3, qwen35, qwen35moe, qwen3next) still pass nine arguments, so only qwen4exp.cpp:1145 changes behaviour; speculative nodes skip add_fused_node; and the flag is off by default, so the shipped path is untouched.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: CHANGES at 5f0d4619ae — the design is careful and I could not find a correctness hole by reading, but the claim this PR makes is a cost/benefit
one and the cost side is not measured.
What is right, and it is most of it. The diagnosis is correct and matches what this model does: with hyper-connections the residual is 4 x 2560 wide, and
l+1's router input is l+1's own mixer applied to the post-MoE residual, so the existing gate-ahead applies the right matrix to the wrong tensor. The
guards on the speculative pass are the part I would have worried about and they are all present: the recurrent state and the conv tail go to scratch
instead of the cache (ssm_override/conv_override), register_fused is false in spec mode so a speculative node cannot enter a fused dispatch's
manifest, the whole thing is getenv("ONEBIT_MOE_PREDICT")-gated with off as the default, and there is a unit test
(tests/test-onebit-moe-predict.cpp). That is a better-defended PR than most.
Why CHANGES anyway. This adds a full extra mixer block per layer per token — attention or DeltaNet plus combine plus the FFN mixer plus the router
matmul — to make a prefetch decision. The streaming doc this PR itself cites records the cheap form of that predictor as a net loss on this drive
(247 s with prefetch vs 133 s without). So the whole PR rests on "a better prediction amortises a much larger predictor", and that is a measurement,
not an argument. The body says "untested on device; host-side build + unit test only" — so the one number that decides it is absent.
What I need before this can be judged:
- The end-to-end row, on the configuration the doc measures: streamed Flash-Next,
ONEBIT_MOE_PREDICT=off|block|routerxdepth1 (default), same
pass, three runs each, pass time in seconds. Your code already keeps a hit-rate table — quote it for the block and router modes, because if the hit
rate does not clear the gate-ahead's ~2/3 then the extra block is paying for noise. - One sentence or an assert on the KV-cell claim. The diff says "exact block rewrites the same cache cell before anything reads it" — that is the
single assumption under which a speculative attention pass is safe, and on this model the full-attention layers write throughqsa_sel/set_rowswith
an index list, so "the same cell" is worth pinning down rather than asserting in a comment. A debug assert that the spec pass's chosen cell set equals
the real layer's would cost nothing on theoffpath. - The
step_end(sched.get())readback. It runs inllama_context::graph_computefor every graph — confirm it is a no-op when no prediction was built
in that graph (the comment says so) and that when a prediction exists it does not add a per-token host sync on an APU, where I have measured
readbacks of this kind to cost milliseconds. If it does, it belongs on a sampled path.
None of this is a redesign; (1) is one claim's worth of measurement, (2) and (3) are small.
Context for the reviewer queue: this is the model my lane has spent the night on (engine ggml-org#679), and its long-context decode is dominated by exactly the
mask/indexer path this PR touches at build_qsa_sel, so a device row here would also be the first data on whether the speculative pass helps the case the
issue is about.
VERDICT: READY as a build-time-flagged change (2 conditions before anyone sets
|
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT at 5f0d461: CHANGES
Re-stated in the form the merge script reads; the reasoning is in the review above and the head has not moved since.
|
VERDICT: READY at The property that makes merging safe is verifiable in the diff and still holds: the feature is off unless asked for. The two conditions from the earlier verdict, checked at this head
Both are recorded here rather than assumed half-done: condition 1 cannot be satisfied by reading code, and condition 2 is a one-line edit someone should |
VERDICT at
|
Untested on device. Host-side build + unit test only; no GPU run was made.
Why
The streamed-expert prefetch on Qwen3.8-Flash-Next paid for a prediction that cannot work on this model. With hyper-connections the residual is 4 x 2560 wide, and layer l+1's router input is l+1's own mixer applied to the residual after layer l's MoE. The gate-ahead form (
ONEBIT_MOE_PREFETCHin the streamer) applies l+1's router to layer l's FFN input, so it mixes through the wrong weights: right about two times in three, and a net loss on a drive that is already the limit (enginedocs/moe-streaming.md, "Flash-Next, streamed": 133 s per pass without prefetch, 247 s with).What this adds
Inside layer l's graph, before l's MoE output exists:
ONEBIT_MOE_PREDICT_DEPTHdeep (deeper than the layer's own top-k, so the near misses are read too), with their logits for the margin gate.Only the ranking is used. Whatever a token really routes to is still waited for, so the model's output is unchanged — the predictor decides which experts a prefetch reads, nothing else.
State-write suppression. The speculative pass must not leave state the exact pass reads: the DeltaNet recurrent state and the conv histories (DeltaNet and PLE) are written into same-shaped scratch tensors, so one scratch per shape serves every layer. An attention successor needs none: the exact block rewrites the same cache cell before anything reads it (
build_attn_qsa's cpy precedes its attention in graph order), and the speculative selector's approximate keys only degrade the prediction. The speculative nodes are also left out of the fused-op manifest (register_fused=false,spec=true), so a fused dispatch cannot claim the wrong node.Files
src/onebit-moe-predict.{h,cpp}— environment, the margin gate, the hit-rate table (per layer: predicted/actual experts used, candidates read, candidates gated), aP/Atrace for offline comparison, and the prefetch sink (set_prefetch).tests/test-onebit-moe-predict.cpp— the host-side policy: margin gate, depth ceiling, floor, layer-mask parsing, effective depth.ctest -R onebitpasses.src/models/qwen4exp.cpp(the layer loop + the predictor itself + optional scratch parameters onbuild_layer_attn_linear/build_pleand aspecflag onbuild_layer_attn/build_qsa_sel/build_layer_attn_linear),src/models/models.h,src/models/delta-net-base.cpp(register_fused, default true),src/llama-graph.cpp(one line: capture the routing each layer really used),src/llama-context.cpp(one line: read the prediction back after the step), CMake for the two new files.Flags
ONEBIT_MOE_PREDICT=off|router|block(off by default),_DEPTH,_MARGIN(read a deeper candidate only while its logit is within that distance of the best),_MIN,_LAYERS(1,2,5-9,47),_NO_ATTN,_OUT(trace),_STATS(per-layer hit rate at exit).routerkeeps the cheap residual-only form, for an A/B on the same build.What is not done
llama.cppruns all 48 layers in one graph, so the readback lands after the step; the consumer has to be either the streamer's in-graph op (reading the prediction as an op input) or a one-graph-per-layer streamed decode. Both need the device. Until then the sink is the trace, which is exactly the hit-rate measurement the device plan asks for: prediction hit rate per layer vs the routing the step really used.-mdis on. The hit ratio is unaffected.Design record
enginedocs/moe-streaming.md, "The next layer's own block": the mechanism, the three state-write cases, the cost per predicted layer, and the two ways to feed the prefetch. The engine side (--moe-predict, docs) PRs separately.🤖 Generated with Claude Code