Cacheable by Design? Training Mixture-of-Experts Routers for Locality Against the Edge Memory-Bandwidth Wall

Summary (Overview)

• Quantified bandwidth wall on edge hardware: Serving Qwen3-235B-A22B (4-bit, 134 GB) on an 8 GB GPU achieves only 0.44 tok/s decode, matching a first-principles bytes-per-token ÷ bandwidth model to within 6%. A naive batching scheme that should amortize disk sweeps collapses at batch size 32 due to paging thrash.

• Novel routing telemetry tool (llama-moe-trace): A zero-model-surgery instrument reveals MoE routing structure with 2.0× temporal locality between adjacent tokens, a 52.5% working set for 95% of traffic, and near-orthogonal per-domain expert sets (code overlaps other domains by only 0.11–0.16 similarity).

• Pre-registered negative result: Training routers with auxiliary locality and domain losses (137M parameters, 200M tokens) reduces cache misses up to 60% and achieves 99% static-pin hit rates, but every configuration fails the strict ≤1% perplexity gate—cacheability and quality are tightly coupled with no loss weight threading the joint bar.

• Complementary stacking result: Training-free cache-aware rerouting stacks with trained locality, reaching ~80% cache-miss reduction at ≤3.4% perplexity—far cheaper than either mechanism alone. This replicates at 340M scale.

• Scale study finding: The locality tax does not shrink with scale (rises slightly from +2.0% at 137M to +2.5% at 340M), though undertraining at matched budgets leaves this suggestive rather than decisive.

Introduction and Theoretical Foundation

Background

Mixture-of-Experts (MoE) architectures decouple total parameter count from active parameters per token via a router that selects a top-k subset of experts [Shazeer et al., 2017, Fedus et al., 2022]. This sparsity enables frontier models like Qwen3-235B-A22B to be "runnable" outside datacenters—235B parameters but only ~22B activated per token. However, on commodity hardware, the key constraint is memory bandwidth, not compute.

The Bandwidth Wall

Autoregressive decode is memory-bandwidth-bound: each token requires streaming active expert weights from whatever memory tier holds them. The central observation is that throughput is governed by:

$$\text{tok/s} = \frac{\text{bandwidth}}{\text{bytes-per-token}}$$

Table 1: Memory hierarchy on the test machine (i9-12900, RTX 3070 8 GB, 32 GB DDR4, WD SN530 NVMe). The 4-bit 235B model is 134 GB.

The 50× gap: achieving a usable 10 tok/s at 12 GB/token requires 120 GB/s sustained bandwidth, while the SSD delivers only 2.4 GB/s.

Theoretical Basis

Every intervention (batching, caching, placement, model changes) attacks either the bytes-per-token term or the bandwidth-to-bytes term. The paper systematically works through both, culminating in an attempt to reduce bytes at training time via router losses designed for cacheability.

Methodology

Serving Measurements

• Model: Qwen3-235B-A22B (Q4 K M, 3-part GGUF, 134 GB) served with llama.cpp
• Streaming: Experts streamed from NVMe SSD on demand
• Measurement: Per-token bandwidth usage (~12 GB/token for ~22B active parameters)

Table 2: Serving measurements — Aggregate throughput at batch B across concurrent streams.

Routing Telemetry (llama-moe-trace)

A ~120-line addition to llama.cpp's eval-callback recording router top-k selections per layer and token, with weights untouched. Critical subtlety: the top-k tensor is a non-contiguous view of an argsort, requiring a stride-honoring copy to avoid garbage data.

Traces collected on Qwen3-30B-A3B (48 layers, 128 experts, top-8) over 8,000 tokens each of prose, code, math, and medical text.

Training Setup

Model Architecture (Table 9):

• 137M parameters (d_model=384, 8 layers, E=16 experts, top-2, expert d_ff=768)
• GPT-2 BPE vocabulary (50257)
• Context length 1024

Data: 300M tokens (later trimmed to 200M), 100M each of:

• Prose (WikiText-103)
• Code (codeparrot-clean)
• Math (OpenWebMath)

Arms:

• Arm A: Balance loss only (baseline)
• Arm B: + $\mathcal{L}_{loc}$ with λ swept (0.02, 0.03, 0.05)
• Arm C: + $\mathcal{L}_{dom}$ with μ=0.1 (domain confinement)

Router Losses

Locality loss (temporal reuse):

$$\mathcal{L}_{loc} = \frac{1}{L}\sum_{\ell} \text{mean}_t\left(1 - \langle p_t^{\ell}, p_{t-1}^{\ell}\rangle\right)$$

Domain confinement loss:

$$\mathcal{L}_{dom} = \frac{1}{L}\sum_{\ell} \text{mean}_t \sum_{e \notin \mathcal{S}(d_t)} p_{t,e}^{\ell}$$

Total loss: $\mathcal{L} = \mathcal{L}_{LM} + \alpha\mathcal{L}_{balance} + \lambda\mathcal{L}_{loc} + \mu\mathcal{L}_{dom}$

Pre-registered Criteria

• RQ1 (locality): ≥30% reduction in LRU miss/token at 25% capacity AND validation perplexity within +1% of baseline
• RQ2 (domain): ≥90% static-pin hit@50% within domain AND perplexity within +2%

Empirical Validation / Results

Routing Structure (Qwen3-30B-A3B)

Table 3: Routing profile

Table 4: Cross-domain expert-usage similarity — Code's expert set is nearly disjoint from other domains.

Main Training Results

Table 5: Main results (200M tokens/arm). Metrics averaged over prose/code/math.

Key findings:

• RQ1 refuted: The locality loss shows clean, monotonic dose-response, but no weight satisfies both gates. λ=0.02 preserves quality (+0.3% PPL) but cuts misses only 15% (<30%); λ=0.05 cuts misses 59–60% but costs +2.1–3.1% PPL (>1%).
• RQ2 refuted on quality: Domain confinement achieves 0.991 static-pin hit rate (clearing 0.90) but at +3.1% perplexity, over the +2% gate.

Training-Free Rerouting Complementarity

Applying inference-time cache-aware rerouting (routing to cached experts within relative tolerance τ) to the locality-trained model is nearly free:

• Baseline: τ=0→0.5 costs +4.0–5.8% PPL
• Trained model: Same intervention costs only +0.5–0.9% PPL (5–8× smaller marginal cost)
• Stacked: Training-time locality (~58–61% miss reduction at +1.4–2.2%) + τ=0.5 reaches 80–82% miss reduction at +1.9–3.1% PPL; τ=0.7 reaches 86–88%

Scale Study

Table 6: Scale rung — The locality tax does not shrink with scale; stacking advantage persists.

Per-Domain Analysis

Table 7: Per-domain metrics — The perplexity tax is concentrated in the hardest domain (prose, +3.0%) and essentially zero for the easiest (code, PPL ≈4.7 both arms), while cacheability gains are consistent everywhere.

Theoretical and Practical Implications

Theoretical Implications

1. The locality tax is real and coupled: Training routers for cacheability at 137M scale incurs a genuine perplexity cost that scales with the strength of the locality constraint. The dose-response curve is tight—no configuration threads the joint miss/quality bar.

2. Scale does not obviously help: The 340M rung suggests the tax does not shrink with scale (+2.0%→+2.5% at fixed ~58% miss reduction), though undertraining at matched budgets leaves this suggestive. This contradicts the optimistic hypothesis from sparsity scaling arguments that larger models have more spare capacity.

3. Contrast with concurrent work: StickyMoE [Kayyam, 2026] reports the same adjacent-token consistency loss as nearly free (even improving perplexity up to −4.1%) on single-domain, sub-25M models. This paper's multi-domain 137M results contradict that—highlighting corpus and scale dependence that reconciles the two findings.

4. Complementarity is the key positive result: Training-time locality co-adapts experts into mutually substitutable neighborhoods, making inference-time substitution within a neighborhood nearly free. This two-stage approach (train for locality + reroute at serving time) reaches ~80% miss reduction at acceptable quality cost.

Practical Implications

1. Path-Mapped Serving (PMS) ceiling: Placement policies can improve decode 2–7× at full accuracy, but single-stream 10 tok/s requires ~95% hit rate needing ~70 GB resident—4× this machine's fast memory. Interactive speed is unreachable on this hardware class.

2. Domain-primed prefetching doesn't help: LRU caches re-warm within O(capacity) tokens of a domain switch, so boundary miss bursts amortize to nothing over realistic ≥1k-token domain runs.

3. Expert elision is possible: On locality-trained models, eliding the second expert when its renormalized gate < 0.02 drops ~12% of expert loads at ≤0.05% perplexity change.

4. Engineering reproducibility: The paper logs six corrections during the study, including a non-contiguous argsort view pitfall, embedding initialization issues, and a 2.8× slower padded-batched-GEMM expert dispatch.

Conclusion

The central finding is a pre-registered negative result: training routers for locality works mechanically (up to 60% cache-miss reduction, 99% static-pin hit rates) but is never quality-free at 137M scale on multi-domain data. The perplexity tax and cacheability gain are tightly coupled, and no loss weight threads the strict joint bar of ≥30% miss reduction at ≤1% perplexity cost.

Key takeaways:

1. The bandwidth wall on edge hardware is real, quantified, and matches first-principles models precisely
2. Routing structure shows exploitable temporal locality (2.0× chance) and domain structure, but this is accidental
3. Training for locality is possible but not free—the tax is real at small scale and does not shrink at 340M
4. The actionable positive result: combining training-time locality with training-free rerouting at serving time reaches ~80% cache-miss reduction at ≤3.4% perplexity—far cheaper than either alone

Future directions:

• A compute-optimal and ≥1B rung remains the definitive test of whether the locality tax shrinks with scale
• Same-artifact routing traces on the actual 235B model (rather than the 30B proxy)
• Purpose-built sequential-sweep batch serving engines to realize the theoretical batching gains

The paper's contribution is the discipline of the pre-registered, strict-criterion, multi-domain evaluation itself—converting a vague hope ("train routers to be cache-friendly") into a precise, falsifiable claim, and reporting the null result in full.