BigMoeOnEdge/docs/architecture.md
Raffaele 9153dcc27a
feat(moe): serve row-gathered dense tables from flash (--row-stream) (#180)
Dense tables the graph only gathers rows from (the token embedding, on
most models) are bound to reserved address space and fetched in 16 KiB
slabs inside a bounded LRU window, instead of being read whole and kept
resident. Which tables qualify is decided from the captured graph, not
from a name list. Byte-identical to the resident reference; -497 MiB
pinned on Qwen3.8-Flash-Next and -515 MiB on Qwen3.6-35B on the 12 GB
test phone, throughput neutral, off by default. Gates G15a/G15b.
App 0.24.0 (unreleased), new page docs/row-gathered-tables.md.
2026-08-29 10:04:02 +02:00

5.6 KiB

Architecture

BigMoeOnEdge is a small ports-and-adapters engine that sits on top of llama.cpp's public API. Its guiding constraint: drive streaming through the public API so upstream updates cost a submodule pointer bump and nothing else. The serial streamer holds to this against stock upstream; the one exception is the optional --overlap feature, which carries a single ~25-line hook on a fork branch with an explicit sunset (see below and seam.md § 3).

Layers

cli/            bmoe-cli — parses flags, the only place env vars are read
core/
  include/bmoe/ ports (interfaces) + config, pure policy, no llama.cpp dependency
    config.h        RunConfig + validate()
    expert_source.h IExpertSource — the residency strategy port
    row_source.h    IRowSource - the row-gathered residency port
    recipe.h        MoeRecipe + registry
    metrics.h       TokenMetrics / RunSummary + IMetricsSink
    runtime.h       run() entry point
  src/
    io/         platform_io — O_DIRECT reads + reserve/commit/evict VM, cross-platform
                file_reader — pooled positioned reader, per-consumer O_DIRECT
    moe/        gguf_offsets (tensor → (shard, offset), split ggufs included), arch_registry,
                expert_stream_source (one reader per shard), router_hook
                dense_weights — non-expert weight policy + the residency sensor
                row_stream - row-gathered tables served from flash (see row-gathered-tables.md)
    engine/     session — composition + the generation loop (open/generate/close)
                runtime — the one-shot run() wrapper over a Session
                chat_parse — reasoning-parser wiring (llama.cpp `common`, see seam.md)
                thinking_control — how "thinking off" is honoured, probed per model
    metrics/    csv_metrics_sink, route_trace_sink, decode_trace_sink
third_party/
  llama.cpp     upstream submodule; public-API consumer, plus one optional overlap hook
tests/          byte-identity gates
examples/android an APK that drives bmoe-cli via ProcessBuilder

Dependencies point inward: adapters depend on the port headers, the CLI composes them. The pure-policy code (config.cpp, arch_registry.cpp) compiles with no native dependency, so a subset of the project builds and is testable before llama.cpp is fetched.

Why the streaming seam needs no fork

Streaming experts serially needs three things from the inference engine. All three are already public in llama.cpp:

  1. A hook at routing time. llama_context_params.cb_eval is called for every graph node. We ask for the routing nodes (ffn_moe_topk-<il>); ggml computes and synchronizes each alone, then calls us back with the selected expert ids materialized. The route trace and cache-aware dropping additionally ask for each layer's ffn_moe_weights*-<il> chain — and dropping is the one path that writes into a graph tensor's contents rather than only rebinding ->data. See seam.md.
  2. The expert tensor pointers. During a one-token warm-up we scan each graph node's sources for tensors named blk.<il>.ffn_{gate,up,down}_exps.weight and record the live ggml_tensor*. We then rebind their ->data.
  3. The file layout. gguf_get_tensor_offset (public) gives each tensor's byte offset so we can pread individual expert slices.

Loading with use_mmap=true, use_extra_bufts=false keeps the weights in their native gguf layout (a repacked buffer would break the rebind). That is a public model parameter.

Because none of this touches llama.cpp internals, the serial streaming path runs against the unmodified upstream repository. Contrast with approaches that patch the model files: those must be rebased on every release. Here, git submodule update --remote and a rebuild is the whole upgrade.

The one place we do carry an extension is the optional --overlap feature. Overlapping a token's expert reads with its expert matmuls needs a per-expert wait point inside the CPU MoE kernel, which no public API exposes, so the submodule pins a fork branch adding one ~25-line readiness hook on top of the upstream commit. It is zero-cost when unregistered, the serial path still builds against stock upstream, and it is dropped the moment upstream ships an equivalent callback. Details and the sunset condition are in seam.md § 3.

See seam.md for the exact callback contract and the ggml behaviour it relies on.

The generation loop

The composition root is Session (core/src/engine/session.cpp):

  1. open() — load model (mmap on, repack off, experts on CPU); if streaming, resolve the architecture recipe, install the router hook, do the capture warm-up, bind the expert source, clear the warm-up KV. Done once per model.
  2. generate() — prefill the prompt, then greedily decode n_predict tokens, reporting per-token metrics. Callable repeatedly; the expert cache stays warm between calls (see session.md). Cancellable mid-flight via the abort callback.
  3. Destructor — tear down in order: I/O pool, context, hook, model, backend.

run() (core/src/engine/runtime.cpp) is a thin one-shot wrapper — open, one generate, close — so the gates and the interactive session share the same code path.

Greedy sampling makes the output a deterministic function of the graph — the property the byte-identity gates assert. That holds with the lossy knobs off. Under --drop-cold-experts the hook edits routing weights from live cache state, which is not in the graph, so output becomes a function of the graph and the run's history.