[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 its layer's slot pool 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 expert view of a demand chunk carries its layer's expert count and input scales beside the slot range (RoutedExperts). The NVFP4 prefill picks its tiles from tokens per expert (nvfp4_moe.prefill_config), and W4A4 quantizes its activations under the max of the input scales, so over a 1,100-slot pool both would read the slot count and every slot's scales, other layers' included; the Triton and W4A4 GEMMs size their tiles by the layer's experts and W4A4 takes the layer's scales, as the stream does.
• 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, its view naming the layer's 160 experts, 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, its tiles sized by the layer's 8 experts, equals the streamed layer's exactly, and no LRU stat is recorded.
• test_triton_prefill_over_cache_slots_sizes_its_tiles_by_the_routed_layer, test_w4a4_prefill_over_cache_slots_takes_the_routed_layers_scales_and_tiles and test_w4a4_prefill_sizes_its_tiles_by_the_routed_experts_and_hands_them_to_its_fallback: over 1,100 slots the Triton and W4A4 kernels align over the slots and hand their GEMMs the layer's 8 experts, and W4A4 takes the layer's input scales, not the slots' (9.0 in every slot against the layer's 2.0 and 3.0); its fallback to Triton passes both on.
• 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_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 new tests fail on both platforms: the layer and cache tests on the missing setting, method or RoutedExperts, and the bit-identity and tile tests on the missing admission and routed_experts; on Linux, where sgl_kernel 0.4.9 is installed, gdevenyi's test also fails at 2100 experts, with no route sorted into its blocks (assert [] == [0, 1, 2, ...]). 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)), and on its head before the tile sizing (87ef7cd), on Linux, it fails at 1,100 slots, where main's dispatch sized the tiles by the 1,100 slots instead of the layer's 8 experts.

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 26.04 (Python 3.14.9, torch 2.14.1, triton 3.8.0, CUDA 13.3, sgl_kernel 0.4.9, flashinfer 0.6.18), over main at ed359ee: 176 passed, 15 skipped (main: 165 passed, 15 skipped).
• Windows 11 (Python 3.14.9, torch 2.14.1, triton-windows 3.8.0, CUDA 13.4) 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): 177 passed, 16 skipped.

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 stays under num_experts and the smallest slot pool. 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. W4A4 is checked by what it is handed, since its kernels need an sm_120 card; Marlin and b12x read the slot ids directly, as in decode, and neither runs here, since Marlin needs vLLM and b12x an sm_120 card.
• The bf16, FP8 and MXFP4 prefills over the slots still size their tiles by the slot count: correct, but their configs differ from the stream's, so only NVFP4 is asserted bit-identical.
• fix(moe): admit a wide offloaded decode step's experts 128 at a time #680 adds the same ENSURE_QUERY_WIDTH and record_stats lines and admits a wide decode step with the same query loop, so whichever of the two lands second keeps one copy of each.

[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.