feat(moe): load a short prefill chunk's experts on demand through the slot cache

[YevheniiKotyrlo]

Problem

With --moe-strategy offload, OffloadMoELayer._prefill_routed streams every expert of each layer for every chunk, however short. On Qwen3.8-Flash-Next 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. #579 has the trace. Decode already loads only the experts its tokens route to, through the slot cache.

Solution

--moe-prefill-demand-tokens N, off (0) by default. A prefill chunk of at most N tokens loads the experts it routes to into the slot cache, as decode does, and runs the same grouped GEMM over the slot cache at their slots, so its output is bit-identical to the streamed layer's. A longer chunk streams as before.

• OffloadMoeCache.loads_prefill_on_demand and a branch at the top of _prefill_routed.
• OffloadMoeCache.load_prefill_experts hands lru_ensure the chunk's distinct experts 128 at a time, the width its [K, K] dedup block stays flat to in flashlib's slot-cache cost model (it spills past that and does not compile past 1024), and copies each query's misses before the next query evicts. The distinct experts are a sorted list of fixed length, so there is no host sync, and the admissions stay out of collect_decode_freq and the LRU stats.
• With sgl_kernel installed, moe_align_block_size aligns on Triton past 1023 experts: sgl_kernel counts experts in one 1024-thread block, and a demand chunk aligns against the slot count.
• The cache refuses the flag unless decode runs on the GPU, and a dense model resets it to 0 with the other MoE settings in _DENSE_MOE_SETTINGS.
• Registered beside the other prefill knobs: EngineConfig, the cache construction, server/args.py, docs/cli.md.

Closes #579.

Tests

• tests/moe/test_offload.py::test_offload_moe_layer_short_prefill_loads_its_experts_through_the_slot_cache: a 96-token chunk routing to 160 experts admits them in queries of 128 and 32, copies each query's misses before the next, counts none of them as decode, and runs the GEMM over the slot cache at their slots without materializing the layer; a 97-token chunk streams.
• test_offload_moe_cache_takes_a_prefill_chunk_on_demand_only_within_its_threshold: a threshold of 0 takes no chunk on demand, and 2 takes chunks of 1 and 2 tokens.
• tests/moe/test_nvfp4_backends.py::test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit (GPU): a 600-token chunk, 1,200 routes, after decode left other experts resident, over a 1,100-slot cache and over a cache of one layer's experts filled three per query. The GEMM over the slots equals the streamed layer's exactly, and no LRU stat is recorded.
• test_adjust_config_drops_the_demand_prefill_threshold_for_a_dense_model: a dense model given 512 ends with 0.
• test_offload_moe_cache_rejects_demand_prefill_without_gpu_decode, cpu and hybrid: each raises needs GPU decode.
• tests/moe/test_fused_moe.py::test_moe_align_leaves_sgl_kernel_past_its_one_block_of_experts: with sgl_kernel installed, 1023 experts align on it, and 1024 and 2100 on Triton.
• test_moe_align_sorts_every_route_into_its_expert_blocks_past_1024_experts (GPU, gdevenyi's): at 1023 and 2100 experts every route lands once, in a block of its expert.

Over main the first six fail: the layer test with no admission (assert [] == [128, 32]), the align test sending 1024 and 2100 experts to sgl_kernel, and the others on the missing setting or method. On this PR's first head (a8f24b1), on Windows, the bit-identity test at 600 tokens fails to compile _lru_ensure_kernel (numel (4194304) exceeds triton maximum tensor numel (1048576)).

A/B

RTX 3090 Ti 24 GB on PCIe 4.0 x16, Ryzen 9 9950X3D, driver 617.14, Windows 11, RadixArk/Qwen3.8-Flash-Next-NVFP4. Two arms of one build, one run each: the flag at 0, which streams every chunk as main does, and at 4,096. Every prompt is a streaming chat completion, timed from the request to its first token:

prompt tokens	87	117	151	230	387	718	1,324
--moe-prefill-demand-tokens 0	2.69 s	2.67	2.64	2.60	2.61	2.60	2.60
--moe-prefill-demand-tokens 4096	1.18	1.36	1.48	1.68	2.05	2.63	3.10

The arms cross near 700 tokens, and I serve with 512. The build is my FreeToken 0.1.3 with my other changes, measured on 2026-09-30; its expert cache admits by LRFU rather than LRU, a whole chunk in one call, so this table does not time the 128-wide admission. gdevenyi timed his versions of the two fixes in this PR on an RTX 5080 16 GB under Linux with a 2,100-slot cache (comment below): at 512, short turns over a cached 28k-token prefix reached their first token 40-70% sooner, longer turns were unchanged, and decode held at 55-58 tok/s.

The server command

ft serve --model RadixArk/Qwen3.8-Flash-Next-NVFP4 --moe-strategy offload --text-model-only \
  --max-running-requests 1 --num-tokens 1048576 --kv-reserve-tokens 1048576 \
  --max-seq-len-override 1048576 --max-output-tokens 1048576 --pin-reasoning-effort xhigh \
  --memory-ratio 0.85 --vram-reserve-mb 2000 --kv-placement host --token-embedding-placement host \
  --moe-cache-policy lrfu --mamba-host-slots 128 --bind-when-ready --moe-collect-stats \
  --hf-overrides '{"text_config":{"rope_parameters":{"rope_type":"yarn","rope_theta":10000000,"partial_rotary_factor":0.25,"mrope_interleaved":true,"mrope_section":[11,11,10],"factor":4,"original_max_position_embeddings":262144}}}' \
  --moe-prefill-demand-tokens 0    # 4096 in the other arm

--vram-reserve-mb is #574 and --kv-placement is #672; --pin-reasoning-effort, --token-embedding-placement, --moe-cache-policy lrfu, --mamba-host-slots, --moe-collect-stats and --bind-when-ready come from my local build.

Verification

The test files that reach the changed code are tests/moe/test_offload.py, test_nvfp4_backends.py, test_fused_moe.py, test_dsfp4_grouped_prefill.py, test_gpt_oss.py, test_prefill_hit_d2d.py, test_hybrid_fetch.py, test_fused_copy.py and tests/engine/test_cache_budget.py.

• Linux, WSL2 Ubuntu 24.04 (Python 3.12, torch 2.14.0, triton 3.8.0), over main at 9c635a1: 125 passed, 7 skipped, and the three tests that need flashinfer, which this environment lacks and which fail on main too.
• Windows 11 over my Windows fixes (fix(windows): build the extensions and JIT kernels with MSVC #575, fix(windows): serve over loopback ZMQ and a selector event loop #580, fix(windows): read checkpoints without O_DIRECT, posix_fadvise or PROT_READ #581): 129 passed, 6 skipped, and two of test_prefill_hit_d2d.py's batch-memcpy tests, which fail now and then until fix(kernel): order the batch-memcpy probe against the current stream #414 lands.
• Over main the first six tests fail on both platforms, as above. Neither environment has sgl_kernel, so the out-of-bounds read itself is gdevenyi's measurement.

The tracker has no other issue or PR that loads a prefill chunk's experts on demand: searched prefill slow, nvfp4 prefill, prefill throughput, fused_nvfp4, prefill demand, experts on demand prefill.

Known limits

• Past about 700 tokens 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 the default stays 0.
• DSV4OffloadMoELayer already takes a chunk on demand while T * top_k < num_experts. On Flash-Next that rule stops at 51 tokens (512 experts at top-10), far below the crossover above, which is why this is a token threshold. The DeepSeek-V4 path is left as it is.
• The bit-identity test runs the Triton NVFP4 kernel. Marlin and b12x read the slot ids directly, as in decode; neither runs here, since Marlin needs vLLM and b12x an sm_120 card.

[gdevenyi]

I tested this on Linux: RTX 5080 16 GB (sm_120), torch 2.14.1, Triton 3.8, sglang-kernel 0.4.9, nvidia/Qwen3.8-Flash-Next-NVFP4 offloaded with a 2100-slot expert cache, merged onto main 9c635a1. Two problems on the demand path at realistic chunk sizes, then the TTFT results with both fixed.

1. ensure_experts gets the chunk's full tokens x top_k ids. The LRU slot-cache kernel's block is next_pow2(ids) x next_pow2(slots). A direct call at 2100 slots, top-10:

chunk tokens (ids)	result
8 (80), 32 (320)	ok (the 32-token call took 33 s)
64 (640)	did not finish in 60 s
128 (1280) and up	CompilationError ... numel (4194304) exceeds triton maximum tensor numel (1048576)

So any chunk between ~64 and 512 tokens fails, which is the range the description recommends. Handing the kernel the chunk's unique ids (at most num_experts) in decode-width pieces (32), copying each piece's misses before the next ensure, and mapping back through torch.unique(..., return_inverse=True) fixes it: 67cea8b, 3ad9d0f. With a scored eviction policy, the earlier pieces' slots also need pinning against the later ones: d7972fb.

2. moe_align_block_size crashes past 1024 experts with sgl_kernel installed. Over the slot cache the grouped GEMM aligns against cache_size, not the layer's 512 experts. sgl_kernel's align reads out of bounds once num_experts + 1 > 1024:

num_experts	sglang-kernel 0.4.9	0.4.5
1022, 1023	ok	ok
1024, 2100	illegal memory access	illegal memory access

The illegal access poisons the CUDA context, so the first short prefill takes the server down. The bf16 and the NVFP4 GEMMs both go through it. The Triton align in kernel/triton/moe_align.py has no such limit: it gives valid results at 2100 and 4096, and output identical to sgl_kernel's below 1024. The fix routes to it past the limit, with a test: 8eb353b. Windows has no sgl_kernel, which is probably why this did not show in your runs.

Results with both fixed. An agent-style conversation over a cached ~28k-token prefix, 8 turns, TTFT per turn by threshold:

new tokens in the turn	0 (stream)	256	512	1024
38-201 (short follow-ups)	1.25-1.32 s	0.42-0.86 s	0.38-0.86 s	0.38-0.87 s
750	1.38 s	1.45 s	1.41 s	1.44 s
1229	1.44 s	1.45 s	1.44 s	1.51 s
2009-2568	1.44-1.53 s	1.42-1.53 s	1.45-1.54 s	1.45-1.62 s

• Short turns get 40-70% lower TTFT. Turns above the threshold are unchanged.
• At 1024, the 750-1229-token turns get slightly slower: such a chunk routes to most of each layer anyway. That is the same crossover you saw near 700 tokens, so 512 is the right setting here too.
• Cold greedy outputs were identical at thresholds 0/256/512/1024, which confirms the bit-identity end to end in the server.
• Decode was unchanged (55-58 tok/s).

[YevheniiKotyrlo]

Thanks for testing it at a real slot count, and for the TTFT run, which the description now cites. 4bea268 hands lru_ensure the chunk's distinct experts 128 at a time, the width flashlib's cost model keeps its dedup flat to, each query's misses copied before the next, with no host sync, and 9ed4e79 aligns on Triton past 1023 experts, with your test; both name you as co-author. Upstream's cache is LRU-only, so the earlier queries need no pin: LRU evicts them last, and a layer never routes to more experts than the cache holds.