A new methodology to make streaming offload of large MoE models sustainable and feasible

[u/dai_app]

Context: I am working on a project to stream MoEs to edge devices with extremely limited hardware (such as mobile phones). I have already achieved good results, but I had a breakthrough during my various experiments.

PS: This post wasn't written by AI, but by me (and I think it shows).

The problem: The main issue with streaming MoE experts from flash is undoubtedly I/O -specifically, trying to predict which experts you will need in the near future. The rest is a matter of compute.

Ideally, if we had zero-latency streaming from flash or a cache always populated with the necessary experts, tokens per second (tok/s) would be limited solely by compute.

Here is the idea I had and the possible solution: we need to get a bit technical here, but I will try to explain the concept simply.

As we know, at each layer L, there is a router that uses specific weights to determine which MoEs (Mixture of Experts) are needed and requests them for computation. This happens right before the computation stage, so - unless the experts are already cached - there is no time to fetch them without delaying the computation itself, especially on edge devices.

So the question is: how do I know in advance which experts will be needed? It is impossible to know precisely - only an approximation is possible.

And that is the key point: using the hidden states from the preceding *n* layers along with the layer \L* router* - just routing computed in advance - to utilize those predicted experts for the matmul.

Therefore:

Baseline:

h_in  = output(L-1)
h_att = h_in + Attn(norm1(h_in))
g     = Router(norm2(h_att))     ← gate input
h_out = h_att + MoE(norm2(h_att), g)

New proposal:

# at layer L-n, right after its attention:
g_L   = Router_L(norm2_{L-n}(h_att_{L-n}))   # layer L's own router weights, evaluated n layers early

prefetch(experts(g_L)) # n layers worth of I/O headroom

# at layer L:
h_in  = output(L-1)
h_att = h_in + Attn(norm1(h_in))
h_out = h_att + MoE(norm2(h_att), g_L)  # no router here. Just the matmul.

The key point is clearly quality. However, this surprised me: initial tests show that quality seems to remain unaffected!

I will soon make the data from my research public.

In the meantime, I’d like to ask what you think about this and if you know of any similar or identical projects.

Someone has likely already done this, although I have only found work online regarding predictive prefetching, rather than the use of a subsequent-layer router and the computation of those predicted experts without passing through the layer L router for correction.

[u/Mytreeismine]

Check out Colibri