Device-native architectures: fine-grained MoE and pre-attention router

· By José Manuel Román
Contents

This post belongs to the series on inference performance in small models. It is the architectural face of a problem we have already looked at from the compute-regime side (the SLM’s inverted roofline) and from the weight-loading side in From disk to HBM. Here the question is different: what if, instead of adapting a large model to the device, we design the model for the device from the first commit?

TL;DR

The default move to get an LLM onto a laptop, a phone or an edge box is to take a dense model built for cloud and compress it: distillation, pruning, quantisation. It is a gesture of reduction: you start from something big and take things away. SmallThinker (arXiv:2507.20984, SJTU IPADS + Zenergize AI) argues for the reverse gesture, designing from scratch, and articulates it in three pieces. First: fine-grained MoE, many small experts with very few activated per token, so that total parameters N (the capacity) decouple from activated parameters A (the compute cost per token). Second: sparse FFN, ReLU-style activation sparsity inside each block, which adds a second level of sparsity on top of the first. Third: a pre-attention router that predicts which experts will be needed before running the attention block and launches the prefetch of those weights from SSD/flash in parallel with attention compute, hiding the storage latency, which is the real bottleneck when the model does not fit entirely in RAM. The authors report SmallThinker-4B-A0.6B and SmallThinker-21B-A3B exceeding ~20 tok/s on consumer CPUs with Q4_0, using ~1 GB and ~8 GB of RAM. The numbers are interesting and the direction is right; the evaluation methodology and the quality cost of activating so little deserve scepticism, and that is what the last part is for.

The analogy: the librarian who anticipates your request

Picture a huge library with a small reading room. You are sitting in the room at a single desk: only a few books fit on it at a time (that is RAM). The bulk of the collection is in the back room, on long shelves that are slow to walk (that is the SSD/flash). And there is a librarian.

The naive method: you read, you reach a point where you need a specific book, you request it, and then the librarian gets up, walks to the back room, looks for it and comes back. Meanwhile, you wait with the page open, making no progress. Every time you need a new book, you pay the full trip to the back room. The reading room spends most of its time waiting, not reading.

The SmallThinker method: the librarian is clever and gets ahead of you. While you are still reading the chapter’s table of contents, working out what it is about and connecting ideas, which in the model is the attention block, they have already glanced over your shoulder, predicted which three or four books you are going to ask for and gone to the back room to fetch them. By the time you finish the contents and formulate the request, the books are already on your desk. You have not waited: the trip to the back room happened in parallel with your reading of the contents.

The analogy holds on four details:

  • The small desk is RAM; the slow back room is the SSD/flash.
  • The books are the MoE’s experts: only a few are on the desk at any moment.
  • Reading the contents is the attention block; requesting and using the books is the FFN/expert block.
  • The librarian who predicts and gets ahead is the pre-attention router: the prediction is made first, and the trip to fetch (the prefetch) overlaps with reading the contents (attention).

The quantitative question running through the whole post is: does the librarian get back in time? The wait is only hidden if the trip to the back room takes less than your reading of the contents. That is the condition t_{\text{attention}} \ge t_{\text{prefetch}}, and we will do it with numbers.

Compressing a dense model vs. designing for device

It is worth laying out the two approaches plainly, because they are not degrees of the same thing: they are different philosophies.

Approach A — compress a dense model built for cloud. You start from, say, a dense 7B–14B model trained to run on an RTX 4090 (24 GB, Ada Lovelace) or on a generic 4×H100 SXM cluster (320 GB, NVLink, native FP8). To fit it into a device you apply three levers, each with its own post: distillation (you train a small student that imitates the teacher), pruning (you remove weights or whole structures) and aggressive quantisation (you go down to 4 bits or fewer). The resulting model is still dense: all its parameters are activated on every token. You have reduced the number of parameters, but the compute pattern is the cloud’s, only smaller.

Approach B — design for device from scratch. Here the device’s constraints enter the architecture, not a later compression phase. The constraints are three and very concrete:

  1. Weak compute. A laptop CPU or a mobile SoC does orders of magnitude fewer FLOPs than a datacenter GPU. This pushes you to minimise the activated parameters per token, not the total ones.
  2. Little RAM. Tens of GB do not fit. This pushes you to keep resident only what is essential and to stream the rest.
  3. Slow storage. The SSD or flash you are forced to stream from has far less bandwidth than a GPU’s HBM. This turns storage I/O into the bottleneck, and pushes you to hide it.

SmallThinker is approach B taken to the detail: each of those three constraints has an architectural answer. Weak compute is attacked with fine-grained MoE + sparse FFN (minimise A). Scarce RAM is attacked with streaming from SSD (resident ≈ A + cache, not N). Slow storage is attacked with the pre-attention router (hide the I/O behind attention). It is no accident that the three pieces fit together: each one solves a constraint, and together they reinforce each other.

An important nuance, so as not to fall into the hype: approach B is neither free nor universally superior. It requires training a new model (you do not reuse existing weights), and the quality ceiling of a model with a very small A is intrinsically bounded, as we will see. The argument is not “B always wins”, but “for the device regime, B attacks the right bottlenecks, and A only attacks them in passing”.

Two levels of sparsity

The central capacity idea is old and well understood in MoE: separate capacity from compute cost. In a MoE, the model has N total parameters spread across experts, but for each token only A parameters are activated (those of the top-k experts the router picks). Compute cost per token scales with A; knowledge capacity scales with N. SmallThinker applies this idea at two superimposed levels.

Level 1 — fine-grained MoE. “Fine-grained” means many small experts instead of few large ones, with very few activated per token. Instead of, say, 8 experts of which you activate 2, you have dozens of experts of which you activate a handful. With smaller experts, the same A is spread across more possible combinations, which gives the router fine granularity and keeps A very low relative to N. The result is an aggressive N/A ratio: a lot of capacity, very little compute per token.

Level 2 — sparse FFN (ReLU-style activation sparsity). This level is orthogonal and operates inside each FFN. With a ReLU-style non-linearity, a large fraction of the intermediate layer’s neurons produce exactly zero for a given token. A neuron that comes out at zero contributes nothing to the output: its matrix-vector multiplication can be skipped. This is activation sparsity: predictable token by token, and exploitable so as not to load or multiply the weight rows/columns corresponding to inactive neurons. It is the same phenomenon exploited by work such as Deja Vu or PowerInfer; SmallThinker builds it in from the factory by choosing activations that favour it.

The combined effect, in one sentence: large N (capacity), tiny A (compute cost per token ≈ proportional to A), and on top of that, inside that A, a fraction of the multiplications is saved through activation sparsity. It is sparsity on top of sparsity.

Classic MoE (coarse-grained)Fine-grained MoE + sparse FFN8 large experts · activates 2large A per expert · coarse granularitymany small experts · activates fewlow total A · fine granularity+ sparse FFN inside each active expert:neurons at 0 (ReLU) →skipped in the computeCapacity = N (all experts) · Cost/token ≈ A (activated) · and inside A, sparse FFN saves moreThe trick: raise N without raising A. Fine granularity allows a far more aggressive N/A ratio.

The pre-attention router: predict and prefetch

Here is the paper’s specific piece, and the one that gives the post its name. The problem it solves is one of I/O scheduling, not of quality.

When the model does not fit entirely in RAM, the expert weights live in SSD/flash and are loaded on demand. The naive flow of a MoE layer is sequential: you run attention, then the router decides which experts are needed, then you load those experts from SSD (waiting), then you run the FFN of those experts. The loading step is pure waiting: the CPU is blocked waiting for bytes from the SSD. In the device regime, where the SSD is slow, that waiting time dominates the decode step.

The pre-attention router breaks the sequentiality by inverting the order of the decision. The observation is that the router does not need this same layer’s attention output to make a reasonable prediction of which experts will be needed: it can predict it from the state it already has before running attention. So:

  1. Before running the layer’s attention block, the router predicts the experts that will be needed.
  2. It launches the prefetch of those experts from SSD/flash asynchronously.
  3. In parallel, the CPU runs the attention block, which is pure compute and does not need the SSD.
  4. When attention finishes, the prefetched experts are (ideally) already in RAM, and the FFN proceeds without waiting.

Storage I/O has been overlapped with attention compute. It is exactly the librarian who goes to the back room while you read the contents.

Naive (sequential): the load from SSD blocksattentionrouterload experts from SSD(blocking wait)expert FFNt_total = t_att + t_load + t_ffnPre-attention router: the prefetch overlaps with attentioncompute thread (CPU)router (pre)attention (t_att)expert FFNI/O thread (SSD)prefetch experts (t_prefetch)The prefetch is hidden ift_att ≥ t_prefetch: by the time attention ends, the experts are in RAM.If t_prefetch > t_att, a wait bubble (t_prefetch − t_att) appears before the FFN. That is the case to avoid.

The hiding condition is the inequality above: the prefetch is completely hidden if and only if

$$t_{\text{attention}} \;\ge\; t_{\text{prefetch}}.$$

If attention takes longer than loading the experts, the load is free (it was already done). If the experts are too large or the SSD too slow, t_prefetch > t_att and a wait bubble equal to t_prefetch − t_att appears. That is why the design needs A to be small (small experts → fewer bytes to prefetch → low t_prefetch) and the grain to be fine: the two things level 1 of sparsity does are not only about saving FLOPs, they are about making the prefetch fit underneath attention.

The maths that matter

Memory footprint: resident N vs. A + cache

The parameter that decides whether the model fits is how much you have to keep resident in RAM at once.

  • Everything in RAM. If you require all the experts to be loaded, the footprint is \approx N (all parameters, multiplied by bytes/parameter according to the quantisation). For a 21B this is prohibitive on a device.
  • Streaming from SSD. If you only keep the active experts resident plus a cache of the recent/likely ones, the footprint drops to \approx A + \text{cache}. The weights that are not in RAM live on SSD and get prefetched when needed. This is where the real saving is: the resident set scales with A, not with N.

The non-expert part of the model (embeddings, attention, router, layernorms) is always resident, but in a fine-grained MoE the bulk of N is in the experts, so the approximation resident ≈ A + cache + dense_part is a good one.

The prefetch calculation, with numbers

Let us put numbers on the analogy. Suppose a consumer SSD at 5 GB/s of sequential read and a quantised expert of size X MB. The time to load one expert is

$$t_{\text{1 expert}} = \frac{X \text{ MB}}{5000 \text{ MB/s}} = \frac{X}{5000}\ \text{s} = \frac{X}{5}\ \text{ms}.$$

Let us make X concrete. In SmallThinker-4B-A0.6B with Q4_0 (~0.5 byte/param effective counting block overhead), a small expert of, say, 4M parameters weighs \approx 4\text{M} \times 0.5 = 2 MB. Loading it costs t_{\text{1 expert}} = 2/5 = 0.4 ms.

Now the scheduling question: if the layer’s attention block takes Y ms, how many experts can I prefetch while attention runs? The number is

$$n_{\text{prefetch}} = \left\lfloor \frac{Y}{t_{\text{1 expert}}} \right\rfloor = \left\lfloor \frac{Y \cdot 5}{X} \right\rfloor.$$

With Y = 2 ms of attention and X = 2 MB per expert: n_{\text{prefetch}} = \lfloor 2 \times 5 / 2 \rfloor = 5 experts. That is, within that layer’s attention window the SSD manages to bring in 5 experts. If the layer’s top-k activates ≤ 5 experts, the prefetch hides them all and t_prefetch ≤ t_att: zero load latency. If the layer needed 8 experts, you would bring 5 for free and pay the load of the remaining 3 as a bubble: (8-5) \times 0.4 = 1.2 ms of waiting per layer. Hence why the design wants fine grain with a small top-k: to fit underneath the attention window.

Two critical observations about this calculation:

  • The 5 GB/s is idealised sequential read. The experts are scattered across the disk; random 4K reads on a consumer SSD are far slower. The effective bandwidth can be a fraction of the nominal one, which reduces n_{\text{prefetch}}. Any methodology reporting tok/s ought to say whether it measures with experts pre-sorted on disk or with realistic access.
  • The attention window Y shrinks with short context and at the start of generation. With short prompts, attention is cheap and may not cover the prefetch; the advantage of the overlap grows with longer sequences. Another detail an honest benchmark should break down.

Weight footprint: why they report ~1 GB for a 4B

Let us do the 4B calculation in Q4_0. Quantisation to 4 bits ≈ 0.5 byte/param, plus a small overhead of per-block scales (Q4_0 adds one FP16 scale every 32 weights, ~0.56 effective byte/param). So:

$$4\text{B} \times 0.5\ \text{B/param} \approx 2\ \text{GB}.$$

That is, the full model in Q4_0 takes up ~2 GB on disk. But the authors report ~1 GB of RAM. A contradiction? No, and understanding why is understanding the design:

  • Not all experts are resident. Only the activated ones (A = 0.6B) plus a cache fit in RAM; the rest lives on SSD and is streamed. 0.6\text{B} \times 0.5 \approx 0.3 GB of active experts, plus the dense part (attention, embeddings, router) and a cache of hot experts.
  • Sparse FFN reduces the work and the useful resident set. The neurons that come out at zero do not need to be materialised for that token.

Adding up active experts + dense part + a reasonable cache, ~1 GB is plausible. But mind the nuance: ~1 GB is the resident set in RAM, not the total footprint in storage, which is still ~2 GB on SSD. Conflating the two, reporting a bare “1 GB”, is misleading if the reader understands “the model takes up 1 GB”. It takes up 2 GB; it keeps 1 GB in RAM. The distinction matters for a device with 2 GB of free storage: there it does not fit.

Analogously, SmallThinker-21B-A3B: 21\text{B} \times 0.5 \approx 10.5 GB on disk; 3\text{B} \times 0.5 \approx 1.5 GB of active experts, and the ~8 GB of RAM reported includes active experts + a generous cache + the dense part. The large cache is what takes it from 1.5 to ~8 GB: you keep many hot experts resident so as not to hit the SSD constantly.

The quality cost: the necessary scepticism

All the machinery above reduces the compute per token to \approx A. But A = 0.6B activated is very little. This is where the enthusiasm needs a brake:

  • Bounded reasoning capacity. A model that activates 0.6B parameters per token has, per token, the compute power of a 0.6B model, not of a 4B one. The total capacity N=4B helps to store more knowledge (more specialised experts), but the processing of each token is still limited by A. For tasks requiring composition and intensive multi-step reasoning, this is a real ceiling, not a detail.
  • The router is a single point of failure for quality. If the fine-grained router picks the wrong experts, and with fine grain there are more decisions to take, quality drops without any speed metric reflecting it. The pre-attention router makes this worse: it predicts the experts before seeing attention, with less information than a post-attention router. The authors ought to report how much quality is lost by predicting early (mismatch between the prefetched expert and the expert a post-attention router would have chosen).
  • The ~20 tok/s needs small print. On which CPU exactly? With what context and generation length (the advantage of the overlap depends on Y)? Cold start included or steady state? Did the SSD have the experts pre-sorted sequentially? An “exceeds 20 tok/s” without those conditions is a marketing number, not a methodological one.
  • A fair comparison. The right question is not “is it fast?”, but “at equal quality on an independent benchmark, is it faster or smaller than an equivalent compressed dense model?”. That requires evals the reader can reproduce, not just tok/s on the authors’ machine.

None of this invalidates the direction. Designing for device is, conceptually, the right approach: it attacks the real bottlenecks (compute, RAM, I/O) in the architecture instead of patching them afterwards. But “20 tok/s in ~1 GB” is a claim about efficiency, and efficiency only means something anchored to a level of quality measured honestly. Until that anchor is clear, the right amount of scepticism is high.

Implications for on-premise and edge inference

  • The SSD becomes part of the inference hierarchy. In cloud, the hierarchy is HBM → RAM. On device, the SSD/flash enters as one more level, and its bandwidth and random-access latency become first-order performance parameters. This connects with From disk to HBM: cold start and weight streaming stop being just a startup problem and become part of the steady state.
  • The heterogeneous edge box starts to make sense. In a mixed environments pattern, a device-native model like SmallThinker runs on the NUC/edge with CPU and SSD, serving locally, while the heavy work stays on the central cluster. The pre-attention router is what makes a GPU-less edge box viable.
  • Capacity planning changes axes. As discussed in Inference capacity planning, on device the resource to plan is not VRAM but the triad resident-RAM / SSD-bandwidth / CPU-FLOPs. A model with small A and overlapped prefetch moves the bottleneck from “does it fit in RAM?” to “does the SSD feed the prefetch in time?”.

Conclusion

SmallThinker is, above all, a change of question. Not “how do I shrink this cloud model so it fits on the device?” but “what would the model look like if I designed it for the device from the first parameter?”. The answer, fine-grained MoE to decouple N from A, sparse FFN to save inside A, and a pre-attention router that hides storage I/O underneath attention, attacks the device’s three constraints (compute, RAM, I/O) in the architecture, not in a later compression phase. The key condition, t_att ≥ t_prefetch, explains why the pieces fit: fine grain does not only save FLOPs, it makes the prefetch fit underneath attention. The reported numbers (~20 tok/s, ~1 GB / ~8 GB of RAM) are promising and the direction is solid; the cost of activating so little and the lack of methodological detail on quality call for caution. Designing for device is the right bet; measuring it honestly is the unfinished business.

See also

  • MoE inference: the call centre with 256 specialists — the conceptual basis of this post: how a router routes tokens to experts and why N and A decouple; read it first if MoE sounds distant.
  • From disk to HBM: cold start and model loading — weight streaming from slow storage, which here stops being a startup problem and becomes steady state via prefetch.
  • Knowledge distillation — the canonical lever of the “compress a cloud dense model” approach, the exact counterpoint to the device-native one.
  • Pruning LLM models — the other reduction lever; useful for comparing “take away from a big one” against “design small from scratch”.
  • Mixed NVIDIA + Intel environments — where a device-native model fits: the edge box with CPU and SSD that serves locally without a GPU.
  • On-premise LLM inference capacity planning — on device the axes to plan are resident RAM, SSD bandwidth and CPU FLOPs, not VRAM.
  • Inverted roofline in small models (sibling of this series, coming soon) — the SLM’s performance regime, which explains why a small A keeps decode memory-bound and where the real ceiling is.
  • Self-speculative decoding with early-exit (sibling of this series, coming soon) — self-spec applied to on-device MoE: how to accelerate decode without an external draft when the model is already small.
  • Aggressive sub-4-bit and ternary quantisation (sibling of this series, coming soon) — Q4_0 and beyond on device: ternary and 2-bit to push the expert footprint on SSD down further.

References