Dynamic host-RAM tiering for large MoE models (Rust experiment)

[meronrudy]

What problem would this solve?

What problem would this solve?

ToshLLM already does unusually good work on the GPU side for AMD Macs: persistent host↔GPU staging, CPU-offloaded MoE, custom Metal attention/GEMM paths, wave64 support, and multi-GPU transfers.

The remaining limitation I’m interested in is placement policy for machines with far more host RAM than VRAM.

A 2019 Mac Pro can have an extreme memory hierarchy — potentially ~1.5 TB host RAM against tens of GB of VRAM per GPU. For large sparse MoE models, most expert weights could remain resident in host RAM indefinitely while only a smaller hot expert set lives in VRAM.

Today, --n-cpu-moe is effectively a static decision made at launch. It cannot respond to which experts are actually hot for the current session, and it cannot use observed router locality to prefetch likely experts.

The question is whether ToshLLM could use host RAM as a managed second-tier model cache instead of treating CPU offload as a fixed split.

Proposed behavior

I’d like to experiment with an optional host-memory tiering runtime that sits above ToshLLM’s existing Metal implementation.

Conceptually:

GGUF / disk
↓
large host-RAM model tier
↓
hot expert / KV working set
↓
GPU VRAM
↓
existing ToshLLM Metal kernels

The runtime would decide what should be resident and what should move next. ToshLLM’s existing Metal code would continue deciding how the copy and computation happen.

I scaffolded a Rust implementation because the host side eventually becomes mostly a residency/state-management problem:

HostOnly
→ PromotionQueued
→ Uploading
→ VramReady
→ InUse
→ Evictable

The proposed sequence is deliberately measurement-first:

1. Read-only MoE tracing
   • record selected expert IDs by layer/token;
   • record VRAM pressure and existing transfer activity;
   • do not move anything.
2. Offline policy simulation
   • compare static placement against LRU/LFU/windowed policies;
   • measure realistic cache hit rates and bytes moved.
3. Single-GPU expert cache
   • only if the traces show useful locality;
   • reuse the existing ggml_metal_*_async staging functions.
4. Predictive prefetch
   • only if expert transitions are actually predictable;
   • initially use simple per-layer transition statistics rather than another ML model.
5. Multi-GPU placement
   • reuse ToshLLM’s existing event-vs-peer-copy behavior rather than implementing another transfer backend.
6. Potential KV host tier
   • keep recent KV in VRAM and colder blocks in host RAM;
   • integrate with the existing KV/prompt-cache work instead of replacing it.

The static path would remain the fallback throughout.

I would consider the experiment worth continuing only if real traces show something roughly like:

• ≥70% expert-cache hit rate at a realistic VRAM budget, or
• ≥25% reduction in host/device bytes on a useful workload.

Otherwise the right outcome is simply to stop.

Alternatives considered

Rewrite more of ToshLLM in Rust: I don’t think this makes sense. Rewriting SwiftUI or the existing Metal backend would touch a lot of working code without addressing the actual bottleneck.

Replace the Metal transfer layer: also unnecessary. ToshLLM already has persistent staging, async host/device copies, multi-GPU event hand-off, peer copies, etc. The proposed runtime should call those existing primitives.

Just increase --n-cpu-moe: this remains the simplest solution, but it is static. It cannot exploit session-level expert locality or dynamically trade expert residency against growing KV requirements.

Immediately implement predictive prefetch: too speculative. I think routing traces should establish whether locality exists first.

Additional context

I’ve scaffolded the experiment in a fork here:

https://github.com/meronrudy/toshllm/tree/rust-tiering-rfc

The scaffold includes:

tosh-tier/
├── tier-core
├── expert-cache
├── trace
├── transfer-sched
├── host-tier
├── kv-tier
├── policy-sim
├── ggml-shim
└── ggml-sys

It also includes an optional Rust↔llama.cpp link shim, deterministic baseline-vs-linked A/B scripts, and a proposed hardware test matrix.

Nothing currently replaces the AMD Metal kernels or performs live expert/KV movement.

I don’t currently have the target Mac Pro hardware, so the main thing I’m asking for before taking the implementation further is help validating the architecture on ToshLLM’s existing AMD test hardware.

The two questions I’d especially appreciate guidance on are:

1. Where is the cleanest point in the current engine to observe top-k MoE expert selections without forcing a GPU synchronization?
2. Would someone with an AMD Mac Pro / large-MoE setup be willing to run the build and baseline A/B scripts before I implement live placement?

If the instrumentation shows no useful expert locality, I’m happy for that to be the end of the experiment.

Before submitting

• I searched the existing issues and did not find the same request.

[engeldlgado]

@meronrudy thanks for writing this up properly, the measurement-first ordering and the stated abort criteria are the right way round and I wish more proposals arrived like this. Before you spend more time on the scaffold, three things about the current engine that change what your traces would mean. All of it is in the tree today, I checked rather than recalled.

--n-cpu-moe is not a placement cache, it is a compute split. It sets a buffer type override so the expert tensors of the first N layers live in a CPU buffer (common/arg.cpp, ggml_backend_cpu_buffer_type). What happens next is decided per op by the scheduler: Metal's offload_op claims a MUL_MAT_ID only when the batch reaches op_offload_min_batch_size, which is 32 by default. Above it the op is offloaded to the GPU and the expert weights are copied up. Below it the op stays on the CPU and the weights are never moved at all. So during decode there is no per-token expert transport to optimise, the experts are computed in place on the host.

Above that threshold, selective caching has nothing to select. The comment on the existing prefetch path says it plainly: with a large batch virtually every expert is used, so the routing ids are not worth waiting for. That is why GGML_SCHED_PREFETCH_EXPERTS uploads whole expert tensors through a second backend instance on the same device, overlapping compute across staging slots, instead of chasing the top-k. If you trace prefill you will see close to 100% of experts touched and you will read it as no locality, when really it is the regime the design already answered.

And the transfers were not the bottleneck they look like. This is the one I would most want you to have before you invest further. Turning on GGML_SCHED_PREFETCH_EXPERTS together with GGML_CPU_NO_REPACK raised MoE-offload prompt processing by 1.8x to 4.4x on this hardware (35B 197 to 470, gemma-4-26B 116 to 515, gpt-oss 57 to 184) at no cost to generation. When we took that apart, the thing that had been slow all along was the AVX2 GEMM over the CPU_REPACK buffer, not the expert uploads. The earlier theory, the same one your proposal starts from, turned out to be wrong. Both are on by default since 0.81.51 whenever ncmoe > 0.

None of that kills the idea, it relocates it. What survives is decode on a machine where host RAM dwarfs VRAM: 8 of 256 experts per layer per token on a 35B-A3B, computed on the CPU today because the batch is below the offload threshold. The open question is whether enough of that working set is stable across tokens to be worth holding in VRAM and computing on the GPU instead. That is a real question and I do not know the answer.

On your first question, and it is better news than you expect. For the layers you care about, during decode, the ids are already in host memory and observing them costs no synchronisation. The op runs on the CPU backend, so the scheduler has already copied every source across, ids included, before it executes. Nothing to add and nothing to wait for. Above the offload threshold the op moves to the GPU and the ids live on device, but that is the prefill regime where the trace would not tell you anything anyway.

So if you want a step 1 that answers your own go/no-go, I would trace decode only, with ncmoe > 0, host side, and ignore prefill entirely. If the hit rate at a realistic VRAM budget clears your 70%, that is worth a conversation about what comes next. If it does not, you will have learned it cheaply and without writing a transfer path.

Two practical notes. Any A/B has to carry GGML_SCHED_PREFETCH_EXPERTS=1 and GGML_CPU_NO_REPACK=1, or the baseline lands well under what the app actually does and everything measures as an improvement over a straw man. And on the Rust side, be aware this project ships with no external dependencies by design, so a second runtime is not something I can take into the tree. A trace harness that lives outside it and answers the locality question is a different matter, and I am happy to point at the exact spots.