diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 97250c2e8856..7993ee3f3031 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -40,6 +40,7 @@ set(LLAMA_CORE_SOURCES llama-quant.cpp llama-sampler.cpp llama-vocab.cpp + onebit-moe-predict.cpp unicode-data.cpp unicode.cpp ) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9dcaf05a6186..46203e742794 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1,5 +1,7 @@ #include "llama-context.h" +#include "onebit-moe-predict.h" + #include "ggml.h" #include "llama-arch.h" #include "llama-graph.h" @@ -2647,6 +2649,10 @@ ggml_status llama_context::graph_compute( LLAMA_LOG_ERROR("%s: ggml_backend_sched_graph_compute_async failed with error %d\n", __func__, status); } + // the next-layer routing predictor's readback + hit-rate table (src/onebit-moe-predict.*); + // a no-op unless ONEBIT_MOE_PREDICT built a prediction in this graph + onebit::moe::step_end(sched.get()); + // fprintf(stderr, "splits: %d\n", ggml_backend_sched_get_n_splits(sched)); return status; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 23e4c90912e2..2219d5e120c1 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1,5 +1,7 @@ #include "llama-graph.h" +#include "onebit-moe-predict.h" + #include "llama-impl.h" #include "llama-model.h" #include "llama-moe-cache.h" @@ -2177,6 +2179,9 @@ ggml_tensor * llm_graph_context::build_moe_ffn( probs = ggml_reshape_3d(ctx0, probs, 1, n_expert, n_tokens); } + // the routing predictor rates its next-layer prediction against this layer's real routing + onebit::moe::graph_actual(ctx0, gf, il, selected_experts); + ggml_tensor * weights = ggml_get_rows(ctx0, probs, selected_experts); // [1, n_expert_used, n_tokens] cb(weights, "ffn_moe_weights", il); diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ad6612647736..7ca51410260b 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -533,7 +533,8 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( ggml_tensor * g, ggml_tensor * b, ggml_tensor * s, - int il) { + int il, + bool register_fused) { const auto * mctx_cur = inp->mctx; const auto kv_head = mctx_cur->get_head(); const uint32_t mem_size = mctx_cur->get_size(); @@ -565,9 +566,11 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( // state s is 4D [S_v, S_v, H_v, n_seqs]; K snapshot slots are written into the output. ggml_tensor * gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, K); - if (n_seq_tokens > 1) { + // a speculative pass is not part of a layer's fused dispatch: no manifest entry, so the node + // runs as a plain op and cannot make the fused claim pick the wrong tensor + if (register_fused && n_seq_tokens > 1) { res->add_fused_node({LLM_FUSED_OP_GDN_CH, gdn_out, il}); - } else { + } else if (register_fused) { res->add_fused_node({LLM_FUSED_OP_GDN_AR, gdn_out, il}); } diff --git a/src/models/models.h b/src/models/models.h index d6ddf9d16778..06d3b7ec8986 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -97,7 +97,8 @@ struct llm_build_delta_net_base : public llm_graph_context { ggml_tensor * g, ggml_tensor * b, ggml_tensor * s, - int il); + int il, + bool register_fused = true); }; struct llm_build_rwkv6_base : public llm_graph_context { @@ -2551,7 +2552,8 @@ struct llama_model_qwen4exp : public llama_model_base { ggml_tensor * cur, ggml_tensor * inp_pos, int * sections, - int il); + int il, + bool spec = false); // dense self-attention over the cells the QSA mask keeps ggml_tensor * build_attn_qsa( @@ -2575,12 +2577,19 @@ struct llama_model_qwen4exp : public llama_model_base { ggml_tensor * inp_pos, ggml_tensor * kq_mask, int * sections, - int il); + int il, + bool spec = false); + // ssm_override/conv_override: write the new state and the conv tail into scratch instead of + // the cache (the routing predictor's speculative pass); spec also skips the fused-node + // manifest entry -- a speculative node is not part of a fused dispatch ggml_tensor * build_layer_attn_linear( llm_graph_input_rs * inp, ggml_tensor * cur, - int il); + int il, + ggml_tensor * ssm_override = nullptr, + ggml_tensor * conv_override = nullptr, + bool spec = false); ggml_tensor * build_layer_ffn( ggml_tensor * cur, @@ -2611,8 +2620,30 @@ struct llama_model_qwen4exp : public llama_model_base { llm_graph_input_rs * inp, ggml_tensor * emb, ggml_tensor * hidden, + int il, + ggml_tensor * conv_override = nullptr); + + // Next-layer routing predictor (src/onebit-moe-predict.*, ONEBIT_MOE_PREDICT): layer il+1's + // own mixer block on this layer's residual, before this layer's MoE output exists, and its + // router ranking. `res_hc` is the residual after this layer's attention is combined. + void build_moe_predict( + llm_graph_input_rs * inp_recr, + llm_graph_input_attn_kv * inp_attn, + const llama_memory_hybrid_idx_context * mctx_hyb, + llm_graph_input_kpool * inp_kpool, + ggml_tensor * ple_emb, + ggml_tensor * res_hc, + ggml_tensor * inp_pos, + int * sections, int il); + // a same-shaped scratch for the speculative pass's state writes: never read, so one per + // shape is enough for every layer + ggml_tensor * build_moe_predict_scratch(ggml_tensor * proto, ggml_tensor ** slot); + ggml_tensor * pred_ssm_ = nullptr; + ggml_tensor * pred_conv_ = nullptr; + ggml_tensor * pred_ple_ = nullptr; + // returns pair of qkv, z std::pair build_qkvz( ggml_tensor * input, diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 88e2e88914c0..921fe665c55a 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -1,4 +1,5 @@ #include "models.h" +#include "onebit-moe-predict.h" #include "llama-impl.h" #include "llama-memory-hybrid-idx.h" #include "llama-memory-recurrent.h" @@ -431,6 +432,10 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inp_pos = build_inp_pos(); ggml_tensor * inp_out_ids = build_inp_out_ids(); + // next-layer routing predictor (src/onebit-moe-predict.*): single-token steps only, and only + // when ONEBIT_MOE_PREDICT asks for it; it allocates this graph's readback tensors + onebit::moe::graph_begin(ctx0, gf, n_tokens == 1 ? n_layer : 0, onebit::moe::depth(n_expert_used), n_expert_used); + ggml_tensor * ple_emb = nullptr; if (hparams.ple_n_heads > 0) { ple_emb = build_inp_ple(mctx_hyb); @@ -481,6 +486,13 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa res_hc = build_hc_combine(res_hc, cur, inject, il); + // predict layer il+1's routing from this residual, before this layer's MoE output exists, + // so its expert reads overlap this layer's compute (ONEBIT_MOE_PREDICT) + if (n_tokens == 1 && il + 1 < n_layer && onebit::moe::wanted(il + 1)) { + build_moe_predict(inp->get_recr(), inp->get_attn(), mctx_hyb, inp_kpool, ple_emb, res_hc, + inp_pos, sections, il); + } + cur = build_hc_mix(res_hc, model.layers[il].hc_ffn_norm, model.layers[il].hc_ffn_down, @@ -741,7 +753,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel( ggml_tensor * inp_pos, ggml_tensor * kq_mask, int * sections, - int il) { + int il, + bool spec) { const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx(); const int64_t idx_dim = hparams.indexer_head_size; @@ -802,7 +815,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel( // which is the lightning indexer with every head weight set to that scale ggml_tensor * weights = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_idx_h, n_tokens), 1.0f/sqrtf((float) idx_dim)); ggml_tensor * score = ggml_lightning_indexer(ctx0, q, pooled, weights, inp_kpool->pool_mask); // [n_pool, n_tokens] - res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, score, il}); + // a speculative pass is not part of a layer's fused dispatch (see build_layer_attn_linear) + if (!spec) { + res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, score, il}); + } cb(score, "indexer_score", il); const int64_t n_top_pool = std::min(n_pool, hparams.indexer_top_k / kpool); @@ -931,14 +947,15 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( ggml_tensor * cur, ggml_tensor * inp_pos, int * sections, - int il) { + int il, + bool spec) { const int64_t n_embd_head = hparams.n_embd_head_v(); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); // indexer reads the same block input as q/k/v; no cache or no ratio means dense const bool qsa = inp_kpool != nullptr && hparams.dsv4_compress_ratios[il] > 0; - ggml_tensor * sel = qsa ? build_qsa_sel(mctx_hyb, inp_kpool, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr; + ggml_tensor * sel = qsa ? build_qsa_sel(mctx_hyb, inp_kpool, cur, inp_pos, inp->get_kq_mask(), sections, il, spec) : nullptr; // Qwen3Next uses a single Q projection that outputs query + gate ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ] @@ -1014,7 +1031,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear( llm_graph_input_rs * inp, ggml_tensor * cur, - int il) { + int il, + ggml_tensor * ssm_override, + ggml_tensor * conv_override, + bool spec) { const auto * mctx_cur = inp->mctx; const int64_t d_inner = hparams.ssm_d_inner; @@ -1054,8 +1074,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear( gate = ggml_reshape_4d(ctx0, gate, 1, num_v_heads, n_seq_tokens, n_seqs); - ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); - ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + // the routing predictor's speculative pass writes this layer's new state and conv tail into + // scratch: the exact block later reads the state this token has to update + ggml_tensor * conv_states_all = conv_override != nullptr ? conv_override : mctx_cur->get_r_l(il); + ggml_tensor * ssm_states_all = ssm_override != nullptr ? ssm_override : mctx_cur->get_s_l(il); ggml_tensor * conv_kernel = model.layers[il].ssm_conv1d; const int64_t conv_kernel_size = conv_kernel->ne[0]; @@ -1120,7 +1142,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear( cb(k_conv, "k_conv_predelta", il); cb(v_conv, "v_conv_predelta", il); - ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il); + ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il, !spec); ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); @@ -1188,6 +1210,100 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_ffn(ggml_tensor * cur, co return cur; } +ggml_tensor * llama_model_qwen4exp::graph::build_moe_predict_scratch(ggml_tensor * proto, ggml_tensor ** slot) { + if (proto == nullptr) { + return nullptr; + } + if (*slot != nullptr) { + const bool same = (*slot)->type == proto->type && + (*slot)->ne[0] == proto->ne[0] && (*slot)->ne[1] == proto->ne[1] && + (*slot)->ne[2] == proto->ne[2] && (*slot)->ne[3] == proto->ne[3]; + // a different geometry: skip the prediction rather than write a state of the wrong shape + return same ? *slot : nullptr; + } + ggml_tensor * t = ggml_new_tensor_4d(ctx0, proto->type, + proto->ne[0], proto->ne[1], proto->ne[2], proto->ne[3]); + cb(t, "moe_predict_scratch", -1); + *slot = t; + return t; +} + +// Predict layer j = il + 1's routing from this layer's residual, which is missing this layer's MoE +// output. The cheap form (mode=router) applies j's FFN mixer and router to it; the full form +// (mode=block) runs j's own mixer block first, because that block's output is the largest term of +// j's router input that a residual-only predictor cannot see. State writes of the speculative +// pass are redirected: the recurrent state and the conv histories go to scratch (the exact block +// reads what this token has to update), and an attention successor needs nothing because the +// exact block rewrites the same cache cell before anything reads it. Nothing here changes the +// output of the model: only which experts a prefetch reads. +void llama_model_qwen4exp::graph::build_moe_predict( + llm_graph_input_rs * inp_recr, + llm_graph_input_attn_kv * inp_attn, + const llama_memory_hybrid_idx_context * mctx_hyb, + llm_graph_input_kpool * inp_kpool, + ggml_tensor * ple_emb, + ggml_tensor * res_hc, + ggml_tensor * inp_pos, + int * sections, + int il) { + const int j = il + 1; + if (j >= n_layer || res_hc == nullptr) { + return; + } + + const auto & cfg = onebit::moe::config(); + const bool block = cfg.mode == onebit::moe::predict_mode::block && + (hparams.is_recr(j) || cfg.attn); + + // j's PLE injection sits before its attention mixer in the layer, and its conv history is + // read-modify-write like the recurrent state: always to scratch, block or not + ggml_tensor * rp = res_hc; + if (hparams.is_ple(j) && ple_emb != nullptr && inp_recr != nullptr) { + rp = build_ple(inp_recr, ple_emb, rp, j, + build_moe_predict_scratch(inp_recr->mctx->get_p_l(j), &pred_ple_)); + } + + if (block) { + ggml_tensor * inj = nullptr; + ggml_tensor * cur = build_hc_mix(rp, + model.layers[j].hc_attn_norm, model.layers[j].hc_attn_down, + model.layers[j].hc_attn_up, model.layers[j].hc_attn_inject, &inj, j); + ggml_tensor * blk = nullptr; + if (hparams.is_recr(j)) { + ggml_tensor * ssm = build_moe_predict_scratch(inp_recr->mctx->get_s_l(j), &pred_ssm_); + ggml_tensor * conv = build_moe_predict_scratch(inp_recr->mctx->get_r_l(j), &pred_conv_); + if (ssm != nullptr && conv != nullptr) { + blk = build_layer_attn_linear(inp_recr, cur, j, ssm, conv, /*spec=*/true); + } + } else { + blk = build_layer_attn(inp_attn, mctx_hyb, inp_kpool, cur, inp_pos, sections, j, /*spec=*/true); + } + if (blk != nullptr) { + rp = build_hc_combine(rp, blk, inj, j); + cb(rp, "moe_predict_res", j); + } + } + + // j's FFN mixer, then its router; only the ranking of the logits is used + ggml_tensor * inj2 = nullptr; + ggml_tensor * cur = build_hc_mix(rp, + model.layers[j].hc_ffn_norm, model.layers[j].hc_ffn_down, + model.layers[j].hc_ffn_up, model.layers[j].hc_ffn_inject, &inj2, j); + ggml_tensor * logits = build_lora_mm(model.layers[j].ffn_gate_inp, cur); // [n_expert, 1] + cb(logits, "moe_predict_logits", j); + + const int k = onebit::moe::depth((int) n_expert_used); + ggml_tensor * ids = ggml_argsort_top_k(ctx0, logits, k); // [k, 1], best first + cb(ids, "moe_predict_ids", j); + // the candidates' own logits, for the margin gate + ggml_tensor * scores = ggml_get_rows(ctx0, + ggml_reshape_3d(ctx0, logits, 1, logits->ne[0], logits->ne[1]), ids); // [1, k, 1] + scores = ggml_reshape_2d(ctx0, scores, k, 1); + cb(scores, "moe_predict_scores", j); + + onebit::moe::graph_prediction(ctx0, gf, j, ids, scores); +} + // PLE n-gram hash embedding: each token gathers ple_n_heads rows of a shared table. // mixed_n = (t[p]*m[0]) ^ ... ^ (t[p-n+1]*m[n-1]); row = mixed_n % vocab[h] + offset[h] // The hash runs host-side because ggml has no int64 and no xor. EOS resets the window. @@ -1369,7 +1485,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_ple( llm_graph_input_rs * inp, ggml_tensor * emb, ggml_tensor * hidden, - int il) { + int il, + ggml_tensor * conv_override) { const int64_t hc = hparams.dsv4_hc_mult; const int64_t hc_dim = hc * n_embd; @@ -1418,7 +1535,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_ple( const int64_t n_seq_tokens = ubatch.n_seq_tokens; // [hist + n_seq_tokens, hc_dim, n_seqs], tokens on ne[0] - ggml_tensor * padded = build_conv_state_at(inp, inp->mctx->get_p_l(il), + ggml_tensor * padded = build_conv_state_at(inp, + conv_override != nullptr ? conv_override : inp->mctx->get_p_l(il), ggml_reshape_3d(ctx0, normalized, hc_dim, n_seq_tokens, n_seqs), hist, hc_dim, il); diff --git a/src/onebit-moe-predict.cpp b/src/onebit-moe-predict.cpp new file mode 100644 index 000000000000..7bd01db6b308 --- /dev/null +++ b/src/onebit-moe-predict.cpp @@ -0,0 +1,372 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "onebit-moe-predict.h" + +#include "ggml-backend.h" +#include "ggml.h" + +#include +#include +#include +#include +#include + +namespace onebit::moe { + +void dump_stats(); + +namespace { + +// "1,2,5-9,47" -> a per-layer flag vector; anything unparsable is ignored, which leaves the +// layer unpredicted rather than aborting a run +std::vector parse_layers(const char * s) { + std::vector out; + if (s == nullptr || *s == '\0') { + return out; + } + const char * p = s; + while (*p != '\0') { + while (*p == ',' || *p == ' ') { + ++p; + } + if (*p < '0' || *p > '9') { + break; + } + char * end = nullptr; + const long a = strtol(p, &end, 10); + if (end == p) { + break; + } + p = end; + long b = a; + if (*p == '-') { + ++p; + b = strtol(p, &end, 10); + if (end == p) { + b = a; + } else { + p = end; + } + } + if (a < 0 || b < a || b > 4095) { + break; + } + if ((long) out.size() <= b) { + out.resize((size_t) b + 1, 0); + } + for (long l = a; l <= b; ++l) { + out[(size_t) l] = 1; + } + while (*p == ' ') { + ++p; + } + } + return out; +} + +predict_config parse() { + predict_config cfg; + + const char * m = getenv("ONEBIT_MOE_PREDICT"); + if (m != nullptr && strcmp(m, "block") == 0) { + cfg.mode = predict_mode::block; + } else if (m != nullptr && strcmp(m, "router") == 0) { + cfg.mode = predict_mode::router; + } else if (m != nullptr && strcmp(m, "off") != 0 && *m != '\0') { + fprintf(stderr, "[moe-predict] unknown ONEBIT_MOE_PREDICT=%s, using off (block|router|off)\n", m); + } + + if (const char * v = getenv("ONEBIT_MOE_PREDICT_DEPTH")) cfg.depth = atoi(v); + if (const char * v = getenv("ONEBIT_MOE_PREDICT_MIN")) cfg.min = atoi(v); + if (const char * v = getenv("ONEBIT_MOE_PREDICT_MARGIN")) cfg.margin = (float) atof(v); + if (const char * v = getenv("ONEBIT_MOE_PREDICT_NO_ATTN")) cfg.attn = atoi(v) == 0; + if (const char * v = getenv("ONEBIT_MOE_PREDICT_LAYERS")) cfg.layers = parse_layers(v); + if (const char * v = getenv("ONEBIT_MOE_PREDICT_OUT")) cfg.trace_path = v; + if (const char * v = getenv("ONEBIT_MOE_PREDICT_STATS")) cfg.stats = atoi(v) != 0; + + if (cfg.depth < 0) cfg.depth = 0; + if (cfg.min < 0) cfg.min = 0; + if (cfg.margin < 0.0f) cfg.margin = 0.0f; + + if (cfg.mode != predict_mode::off) { + size_t n_sel = 0; + for (uint8_t v : cfg.layers) n_sel += v != 0; + fprintf(stderr, "[moe-predict] %s, depth %d, margin %.2f, min %d, %zu layer(s)%s%s\n", + cfg.mode == predict_mode::block ? "block" : "router (no block)", + cfg.depth, cfg.margin, cfg.min, cfg.layers.empty() ? (size_t) 0 : n_sel, + cfg.attn ? "" : ", attention successors skipped", + cfg.trace_path.empty() ? "" : (", trace " + cfg.trace_path).c_str()); + } + return cfg; +} + +const predict_config & cfg() { + static const predict_config c = parse(); + static const bool at_exit = [] { + if (c.stats) { + std::atexit([] { dump_stats(); }); + } + return true; + }(); + (void) at_exit; + return c; +} + +// --- this graph's readback tensors ------------------------------------------------------------ +ggml_tensor * t_pred_ids = nullptr; // I32 [depth, n_layer] +ggml_tensor * t_pred_score = nullptr; // F32 [depth, n_layer] +ggml_tensor * t_act_ids = nullptr; // I32 [n_used, n_layer] +std::vector has_pred; // per layer: a prediction was copied this build +int g_depth = 0; +int g_n_used = 0; + +std::vector g_layer; +totals_t g_totals; +std::vector g_ids; +std::vector g_scores; +std::vector g_act; + +prefetch_fn g_sink = nullptr; +void * g_sink_ud = nullptr; + +FILE * g_trace = nullptr; +bool g_trace_open_tried = false; + +void trace_open() { + if (g_trace_open_tried) { + return; + } + g_trace_open_tried = true; + if (cfg().trace_path.empty()) { + return; + } + g_trace = fopen(cfg().trace_path.c_str(), "a"); + if (g_trace == nullptr) { + fprintf(stderr, "[moe-predict] cannot open %s for the trace\n", cfg().trace_path.c_str()); + } +} + +void ensure_layer_stats(int n_layer) { + if ((int) g_layer.size() < n_layer) { + g_layer.resize((size_t) n_layer); + } +} + +} // namespace + +const predict_config & config() { return cfg(); } + +bool enabled() { return cfg().mode != predict_mode::off; } + +bool wanted(int il) { + const auto & c = cfg(); + if (c.mode == predict_mode::off || il < 0) { + return false; + } + return c.layers.empty() || ((size_t) il < c.layers.size() && c.layers[(size_t) il] != 0); +} + +int depth(int n_expert_used) { + const int d = cfg().depth > 0 ? cfg().depth : n_expert_used; + return std::max(1, d); +} + +int keep_count(const predict_config & c, const float * scores, int n, int keep_min) { + if (n <= 0 || scores == nullptr) { + return 0; + } + const int want = c.depth > 0 ? std::min(c.depth, n) : n; + // a candidate beyond the depth is read only while it sits close enough to the best one to + // plausibly be the expert the true routing swapped in; the margin gate is what keeps a deeper + // prediction from costing proportionally more bytes. A wrong prefetch still only costs + // bandwidth, so the gate only trims a tail it is sure about. + int keep = want; + if (c.margin > 0.0f) { + const float best = scores[0]; + for (int i = 1; i < want; ++i) { + if (best - scores[i] > c.margin) { + keep = i; + break; + } + } + } + keep = std::max(keep, std::min(keep_min, want)); + return keep; +} + +const std::vector & layer_stats() { return g_layer; } +const totals_t & totals() { return g_totals; } + +void reset_stats() { + std::fill(g_layer.begin(), g_layer.end(), layer_stat{}); + g_totals = totals_t{}; +} + +void dump_stats() { + if (g_totals.asked == 0) { + return; + } + fprintf(stderr, "[moe-predict] %llu step(s), %llu prediction(s): %llu/%llu experts hit (%.1f%%), " + "%llu/%llu candidates read (%.1f%% gated)\n", + (unsigned long long) g_totals.steps, (unsigned long long) g_totals.asked, + (unsigned long long) g_totals.hits, (unsigned long long) g_totals.actual, + g_totals.actual ? 100.0 * (double) g_totals.hits / (double) g_totals.actual : 0.0, + (unsigned long long) g_totals.read, (unsigned long long) g_totals.pred, + g_totals.pred ? 100.0 * (double) g_totals.gated / (double) g_totals.pred : 0.0); + for (size_t il = 0; il < g_layer.size(); ++il) { + const auto & s = g_layer[il]; + if (s.asked == 0) { + continue; + } + fprintf(stderr, "[moe-predict] layer %3zu: hit %5.1f%% (%llu/%llu), read %.1f/layer, gated %llu\n", + il, s.actual ? 100.0 * (double) s.hits / (double) s.actual : 0.0, + (unsigned long long) s.hits, (unsigned long long) s.actual, + (double) s.read / (double) s.asked, (unsigned long long) s.gated); + } +} + +bool graph_begin(ggml_context * ctx0, ggml_cgraph * gf, int n_layer, int n_depth, int n_used) { + if (!enabled() || ctx0 == nullptr || gf == nullptr || n_layer <= 0) { + t_pred_ids = t_pred_score = t_act_ids = nullptr; + return false; + } + g_depth = std::max(1, n_depth); + g_n_used = std::max(1, n_used); + has_pred.assign((size_t) n_layer, 0); + ensure_layer_stats(n_layer); + + t_pred_ids = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, g_depth, n_layer); + t_pred_score = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, g_depth, n_layer); + t_act_ids = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, g_n_used, n_layer); + ggml_set_output(t_pred_ids); + ggml_set_output(t_pred_score); + ggml_set_output(t_act_ids); + // keep them allocated even when no layer writes them (a graph with the predictor off-handed) + ggml_build_forward_expand(gf, t_pred_ids); + ggml_build_forward_expand(gf, t_pred_score); + ggml_build_forward_expand(gf, t_act_ids); + return true; +} + +void graph_prediction(ggml_context * ctx0, ggml_cgraph * gf, int il, ggml_tensor * ids, ggml_tensor * scores) { + if (t_pred_ids == nullptr || ids == nullptr || scores == nullptr) { + return; + } + const int64_t depth_ = ids->ne[0]; + if (depth_ <= 0 || ids->ne[1] != 1 || il < 0 || (size_t) il >= has_pred.size() || depth_ > g_depth) { + return; + } + ggml_tensor * dst_ids = ggml_view_1d(ctx0, t_pred_ids, depth_, (size_t) il * t_pred_ids->nb[1]); + ggml_tensor * dst_sco = ggml_view_1d(ctx0, t_pred_score, depth_, (size_t) il * t_pred_score->nb[1]); + ggml_tensor * ids2 = ggml_is_contiguous(ids) ? ids : ggml_cont(ctx0, ids); + ggml_tensor * sco2 = ggml_is_contiguous(scores) ? scores : ggml_cont(ctx0, scores); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, ids2, dst_ids)); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, sco2, dst_sco)); + has_pred[(size_t) il] = 1; +} + +void graph_actual(ggml_context * ctx0, ggml_cgraph * gf, int il, ggml_tensor * ids) { + if (t_act_ids == nullptr || ids == nullptr) { + return; + } + // the routing is per token and the predictor only runs for a single-token step + if (ids->ne[0] != g_n_used || ids->ne[1] != 1 || il < 0 || il >= t_act_ids->ne[1]) { + return; + } + ggml_tensor * dst = ggml_view_1d(ctx0, t_act_ids, g_n_used, (size_t) il * t_act_ids->nb[1]); + ggml_tensor * src = ggml_is_contiguous(ids) ? ids : ggml_cont(ctx0, ids); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, src, dst)); +} + +void set_prefetch(prefetch_fn fn, void * ud) { + g_sink = fn; + g_sink_ud = ud; +} + +void step_end(ggml_backend_sched * sched) { + if (t_pred_ids == nullptr || sched == nullptr || has_pred.empty()) { + return; + } + ggml_backend_t b_ids = ggml_backend_sched_get_tensor_backend(sched, t_pred_ids); + ggml_backend_t b_sco = ggml_backend_sched_get_tensor_backend(sched, t_pred_score); + ggml_backend_t b_act = ggml_backend_sched_get_tensor_backend(sched, t_act_ids); + if (b_ids == nullptr || b_sco == nullptr || b_act == nullptr) { + // not scheduled in the graph that just ran (a graph built without the predictor) + return; + } + + const int64_t n_layer = t_pred_ids->ne[1]; + g_ids.assign((size_t) g_depth, 0); + g_scores.assign((size_t) g_depth, 0.0f); + g_act.assign((size_t) g_n_used, 0); + + const auto & c = cfg(); + trace_open(); + + g_totals.steps++; + const uint64_t step = g_totals.steps - 1; + + for (int64_t il = 0; il < n_layer; ++il) { + if ((size_t) il >= has_pred.size() || has_pred[(size_t) il] == 0) { + continue; + } + ggml_backend_tensor_get(t_pred_ids, g_ids.data(), (size_t) il * t_pred_ids->nb[1], (size_t) g_depth * sizeof(int32_t)); + ggml_backend_tensor_get(t_pred_score, g_scores.data(), (size_t) il * t_pred_score->nb[1], (size_t) g_depth * sizeof(float)); + ggml_backend_tensor_get(t_act_ids, g_act.data(), (size_t) il * t_act_ids->nb[1], (size_t) g_n_used * sizeof(int32_t)); + + uint64_t hits = 0; + for (int i = 0; i < g_depth; ++i) { + for (int k = 0; k < g_n_used; ++k) { + if (g_ids[(size_t) i] == g_act[(size_t) k]) { + hits++; + break; + } + } + } + const int keep = keep_count(c, g_scores.data(), g_depth, c.min > 0 ? c.min : g_n_used); + + auto & st = g_layer[(size_t) il]; + st.asked++; st.hits += hits; st.pred += (uint64_t) g_depth; + st.actual += (uint64_t) g_n_used; st.read += (uint64_t) keep; + st.gated += (uint64_t) (g_depth - keep); + + g_totals.asked++; g_totals.hits += hits; g_totals.pred += (uint64_t) g_depth; + g_totals.actual += (uint64_t) g_n_used; g_totals.read += (uint64_t) keep; + g_totals.gated += (uint64_t) (g_depth - keep); + + if (g_trace != nullptr) { + fprintf(g_trace, "P %llu %lld %d", (unsigned long long) step, (long long) il, keep); + for (int i = 0; i < keep; ++i) { + fprintf(g_trace, " %d:%.6g", g_ids[(size_t) i], (double) g_scores[(size_t) i]); + } + fprintf(g_trace, "\n"); + fprintf(g_trace, "A %llu %lld %d", (unsigned long long) step, (long long) il, g_n_used); + for (int k = 0; k < g_n_used; ++k) { + fprintf(g_trace, " %d", g_act[(size_t) k]); + } + fprintf(g_trace, "\n"); + } + + if (g_sink != nullptr) { + g_sink(g_sink_ud, (int) il, g_ids.data(), keep); + } + } + if (g_trace != nullptr) { + fflush(g_trace); + } +} + +} // namespace onebit::moe diff --git a/src/onebit-moe-predict.h b/src/onebit-moe-predict.h new file mode 100644 index 000000000000..d9714d91c8bf --- /dev/null +++ b/src/onebit-moe-predict.h @@ -0,0 +1,123 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +// Next-layer-block routing predictor for a streamed MoE (docs/moe-streaming.md, "Next"). +// +// A model whose routed experts come from the drive pays one full layer of reads per layer. The +// prefetch that hides them needs layer L+1's experts while layer L computes, i.e. before L's own +// MoE output exists. The cheap predictor -- L+1's router applied to L's FFN input -- is wrong on +// this family: with hyper-connections L+1's router input is L+1's own mixing step of the residual +// *after* L's MoE, so the mixing step is missing and two errors stack. +// +// So the graph predicts through L+1's own token-mixer block: its PLE injection, its attention or +// DeltaNet block, its combine, its FFN mixer, its router. State writes are suppressed where a +// later exact read would see them (the recurrent state and the conv histories); an attention +// successor needs none, because the exact block rewrites the same cache cell before anything +// reads it. 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. +// +// Environment (read once): +// ONEBIT_MOE_PREDICT=off|router|block off (default), router only (no block, cheap), block +// ONEBIT_MOE_PREDICT_DEPTH=N candidates kept per layer (default: the layer's top-k) +// ONEBIT_MOE_PREDICT_MARGIN=F read a tail candidate only while its logit is within F +// of the best one (default 0: read every kept candidate) +// ONEBIT_MOE_PREDICT_MIN=N never read fewer than N candidates per layer +// ONEBIT_MOE_PREDICT_LAYERS=1,2,5-9,47 predict only these successor layers (default: all) +// ONEBIT_MOE_PREDICT_NO_ATTN=1 skip the block for sparse-attention successors +// ONEBIT_MOE_PREDICT_OUT=FILE append one P/A line per predicted layer per step +// ONEBIT_MOE_PREDICT_STATS=1 one hit-rate line per layer when the context closes +// +// The prediction is emitted as a graph output tensor, read back on the host after the step, rated +// against the routing the step actually used, and handed to whatever prefetch sink is installed. +// A streamer's in-graph op registers itself with set_prefetch(); with nothing installed the sink +// is the trace, so the hit rate is measurable before the streamer exists. + +#include +#include +#include + +struct ggml_context; +struct ggml_cgraph; +struct ggml_backend_sched; +struct ggml_tensor; + +namespace onebit::moe { + +enum class predict_mode { + off, + router, + block, +}; + +struct predict_config { + predict_mode mode = predict_mode::off; + int depth = 0; // 0: the layer's own top-k + int min = 0; // 0: depth + float margin = 0.0f; // 0: no gate + bool attn = true; // block for sparse-attention successors too + std::vector layers; // empty: every layer + std::string trace_path; + bool stats = false; +}; + +// Parsed from the environment on first use; immutable afterwards. +const predict_config & config(); + +bool enabled(); // mode != off +bool wanted(int il_successor); // predict this successor layer +int depth(int n_expert_used); // candidates kept for a layer that routes to n_expert_used + +// The candidate policy, pure host code (tests/test-onebit-moe-predict.cpp). +// `scores` are the candidate logits, best first; the ids run in the same order, so the answer is +// how many of the leading candidates to read: min(depth, n), minus a tail whose margin to the best +// exceeds the configured one, and never fewer than `keep_min`. +int keep_count(const predict_config & cfg, const float * scores, int n, int keep_min); + +struct layer_stat { + uint64_t asked = 0; // steps this layer was predicted + uint64_t hits = 0; // predicted candidates the layer really used + uint64_t pred = 0; // predicted candidates + uint64_t actual = 0; // experts the layer really used + uint64_t read = 0; // candidates the gate would read + uint64_t gated = 0; // candidates the gate dropped +}; + +struct totals_t { + uint64_t steps = 0, asked = 0, hits = 0, pred = 0, actual = 0, read = 0, gated = 0; +}; + +const std::vector & layer_stats(); +const totals_t & totals(); +void reset_stats(); +void dump_stats(); // per layer + totals, to stderr + +// --- graph side ------------------------------------------------------------------------------- +// Called once when the model graph is built with the predictor on: allocates this graph's +// prediction/actual readback tensors. False when the predictor is off. +bool graph_begin(ggml_context * ctx0, ggml_cgraph * gf, int n_layer, int depth, int n_used); +// Per layer: copy the prediction (ids best first, and their logits) into column `il`. +void graph_prediction(ggml_context * ctx0, ggml_cgraph * gf, int il, ggml_tensor * ids, ggml_tensor * scores); +// Per layer: copy the routing the layer actually used into column `il`. +void graph_actual(ggml_context * ctx0, ggml_cgraph * gf, int il, ggml_tensor * ids); +// Called after the step's compute (llama_context::graph_compute): read the tensors back, rate the +// step, write the trace, and hand each layer's gated candidate set to the prefetch sink. +void step_end(ggml_backend_sched * sched); + +// The prefetch consumer (the streamer's in-graph op). `ids` are valid only during the call. +using prefetch_fn = void (*)(void * ud, int il, const int32_t * ids, int n); +void set_prefetch(prefetch_fn fn, void * ud); + +} // namespace onebit::moe diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c44189cebd93..d9456d602ecf 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -148,6 +148,9 @@ if (LLAMA_LLGUIDANCE) llama_build_and_test(test-grammar-llguidance.cpp ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-llama-bpe.gguf) endif () +llama_build(test-onebit-moe-predict.cpp) +target_include_directories(test-onebit-moe-predict PRIVATE ${PROJECT_SOURCE_DIR}/src) +llama_test(test-onebit-moe-predict) llama_build(test-recurrent-state-rollback.cpp) llama_build(test-save-load-state.cpp) diff --git a/tests/test-onebit-moe-predict.cpp b/tests/test-onebit-moe-predict.cpp new file mode 100644 index 000000000000..88bdf3a6fdf9 --- /dev/null +++ b/tests/test-onebit-moe-predict.cpp @@ -0,0 +1,109 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// The routing predictor's host-side policy (src/onebit-moe-predict.*): environment parsing and +// the margin gate. No model, no graph, no device -- the numbers the device plan turns on, checked +// on any machine. + +#include "onebit-moe-predict.h" + +#include +#include +#include +#include +#include + +using namespace onebit::moe; + +static bool near(float a, float b) { return fabsf(a - b) < 1e-5f; } + +// asserts vanish in a Release build, and a test that checks nothing is worse than no test +#define CHECK(x) do { if (!(x)) { fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #x); return 1; } } while (0) + +int main() { + // ---- the margin gate: a deeper prediction is read only where it is close to the best candidate + { + predict_config c; + c.depth = 20; + c.margin = 0.0f; + const std::vector s(20, 0.0f); + // no gate: everything kept + CHECK(keep_count(c, s.data(), 20, 10) == 20); + + // gate 1.0 from the best: candidates at 0, -0.5 kept, -1.5 and beyond dropped + std::vector g(20, 0.0f); + g[0] = 5.0f; g[1] = 4.8f; g[2] = 4.4f; g[3] = 3.9f; g[4] = 3.8f; + for (size_t i = 5; i < g.size(); ++i) g[i] = 1.0f; + c.margin = 1.0f; + CHECK(keep_count(c, g.data(), 20, 10) == 10); // dropped to the floor, never below it + CHECK(keep_count(c, g.data(), 20, 3) == 3); // 5.0 - 3.9 = 1.1 > 1.0 stops before i = 3 + c.margin = 1.5f; + CHECK(keep_count(c, g.data(), 20, 3) == 5); // 5.0 - 3.9 = 1.1 <= 1.5, 5.0 - 3.8 = 1.2 <= 1.5 + c.margin = 0.05f; + CHECK(keep_count(c, g.data(), 20, 1) == 1); // 5.0 - 4.8 = 0.2 > 0.05 + + // the depth is the ceiling whatever the margin + c.depth = 3; + c.margin = 100.0f; + CHECK(keep_count(c, g.data(), 20, 2) == 3); + // fewer candidates than the depth + CHECK(keep_count(c, g.data(), 2, 2) == 2); + // degenerate inputs + CHECK(keep_count(c, nullptr, 20, 2) == 0); + CHECK(keep_count(c, g.data(), 0, 2) == 0); + printf("ok margin gate / depth / floor\n"); + } + + // ---- environment: parsed once, immutable, and the layer mask bounds the prediction + // (set before the first call into the module; config() caches) + setenv("ONEBIT_MOE_PREDICT", "block", 1); + setenv("ONEBIT_MOE_PREDICT_DEPTH", "16", 1); + setenv("ONEBIT_MOE_PREDICT_MIN", "8", 1); + setenv("ONEBIT_MOE_PREDICT_MARGIN", "0.75", 1); + setenv("ONEBIT_MOE_PREDICT_LAYERS", "1,2,5-9,47", 1); + setenv("ONEBIT_MOE_PREDICT_NO_ATTN", "1", 1); + setenv("ONEBIT_MOE_PREDICT_STATS", "0", 1); + + const predict_config & c = config(); + CHECK(enabled()); + CHECK(c.mode == predict_mode::block); + CHECK(c.depth == 16); + CHECK(c.min == 8); + CHECK(near(c.margin, 0.75f)); + CHECK(!c.attn); + + CHECK(wanted(1) && wanted(2) && wanted(5) && wanted(9) && wanted(47)); + CHECK(!wanted(0) && !wanted(3) && !wanted(10) && !wanted(-1)); + + CHECK(depth(10) == 16); // an explicit depth wins over the layer's top-k + CHECK(depth(32) == 16); + + // a gate with the configured floor + { + std::vector s(16); + for (int i = 0; i < 16; ++i) s[(size_t) i] = 10.0f - (float) i; // 10, 9, 8, ... + // best 10: i = 8 gives a margin of 8 > 0.75, so the tail stops there, but min = 8 holds it + CHECK(keep_count(c, s.data(), 16, c.min) == 8); + } + + // ---- unknown mode falls back to off instead of predicting with a mode nobody defined + reset_stats(); + const totals_t & t = totals(); + CHECK(t.asked == 0 && t.steps == 0); + + printf("ok ONEBIT_MOE_PREDICT=block, depth 16, margin 0.75, min 8, layers 1,2,5-9,47, no attn\n"); + printf("test-onebit-moe-predict: all checks passed\n"); + return 0; +}