An offloaded prefill streams every expert of each layer, however short the chunk #579

Before you start

• I have read the Roadmap and this is not already planned there.
• I have read the FAQ.
• I have searched existing issues and found no duplicate.

Applies to

Engine

What problem does this solve

With --moe-strategy offload, a prefill chunk streams every expert of each layer into the double buffer, however short the chunk. On Qwen3.8-Flash-Next (48 MoE layers of 512 NVFP4 experts) that is 1.42 GB a layer and 63.5 GiB a chunk, about 2.6 s over PCIe 4.0 x16 before the first token of any prompt. Traced under Nsight Systems on an RTX 3090 Ti:

In a long conversation that is every turn: the prefix comes from the cache, and the new message still streams all 63.5 GiB.

Proposed solution

A short chunk routes to a small part of each layer, and decode already loads exactly that part through the slot cache (ensure_experts then copy_missing). I tried the same for prefill: a chunk of at most N tokens admits its routed experts into the slot cache and runs the same grouped GEMM over the slots, so its output is bit-identical to the streamed layer's. Every chunk of up to 4,096 tokens on demand against every chunk streamed, time to first token:

The two cross near 700 tokens. Past that, a chunk routes to most of each layer, and on demand its copies wait on the router instead of running behind the previous layer. So I made it a threshold that is off by default, --moe-prefill-demand-tokens N, and I run it at 512. It needs GPU decode; a cpu or hybrid decode target refuses the flag. These numbers come from my build, whose expert cache admits by LRFU rather than LRU. The demand path calls the same ensure_experts either way.

The change is small:

• the flag;
• OffloadMoeCache.loads_prefill_on_demand;
• a branch at the top of OffloadMoELayer._prefill_routed;
• ensure_experts(..., counts_as_decode=False), which keeps prefill admissions out of the decode miss statistics.

The branch is https://github.com/YevheniiKotyrlo/FreeToken/tree/feat-moe-prefill-demand. Its tests check that the layer sends a short chunk through the slot cache and a longer one through the stream, and that the GEMM over the slots equals the stream bit for bit. Would you take it as a PR?

Alternatives considered

• --moe-prefill-hit-d2d copies the experts the cache already holds device-side and only the rest over PCIe, but still fills the whole layer's buffer. On top of this change (with fix(kernel): order the batch-memcpy probe against the current stream #414's probe fix), it moved only a 718-token prompt out of the run-to-run spread.
• A smaller --max-prefill-length does not help: a chunk of any length streams the whole layer.