fix(moe): admit a wide offloaded decode step's experts 128 at a time

[YevheniiKotyrlo]

Problem

With --moe-strategy offload and GPU decode, OffloadMoELayer._decode_routed hands flashlib's lru_ensure every route of a decode step in one call: at the largest captured batch, --max-running-requests times top_k ids. The kernel dedups its query in a [K, K] register block, which flashlib's slot-cache cost model (flashlib/kernels/slot_cache/cost.py, fitted to H100 runs) keeps flat to K = 128 and measures spilling past it. Past 1,024 ids the block exceeds Triton's tensor limit and the kernel does not compile:

triton.compiler.errors.CompilationError: at 10:64:
def _lru_ensure_kernel(
...
numel (4194304) exceeds triton maximum tensor numel (1048576)

At top-10 (Qwen3.8-Flash-Next) that is from 103 concurrent requests. Below it the spill is the cost: on an RTX 3090 Ti one admission takes 7-17 µs a layer up to 12 requests, 44 µs from 13 to 24, and 264 µs at 32 and 48, about 13 ms of every decode step across 48 MoE layers (Measurement below). DeepSeek-V4's on-demand prefill chunks (DSV4OffloadMoELayer._prefill_routed, taken while T * top_k < num_experts) reach 252 routes at 256 experts. #676 hit the same limit on a demand prefill chunk, where gdevenyi found it.

Solution

OffloadMoeCache.load_experts(layer_id, topk_ids) replaces ensure_experts plus copy_missing at both call sites and returns the slot ids:

• A step of at most ENSURE_ONE_CALL_WIDTH (256) routes is the single admission it was, rewriting topk_ids in place: up to 25 requests at top-10 nothing changes, and up to 256 ids the spill costs less than splitting the query.
• A wider step goes to lru_ensure as its distinct experts, ENSURE_QUERY_WIDTH (128) at a time, each query's misses copied before the next query evicts. The distinct experts are a sorted list of fixed length, so the step needs no host sync and stays CUDA-graph capturable. LRU evicts an earlier query's slots last, and a layer never routes to more experts than the cache holds, so no query evicts an earlier one's.
• The statistics count a wide step once, as one admission would: its distinct experts, its misses and one call. decode_freq counts every route.

ENSURE_QUERY_WIDTH and ensure_experts's record_stats parameter are the same lines as in #676, so whichever of the two lands second rebases cleanly.

Measurement

RTX 3090 Ti 24 GB, Windows 11, torch 2.14.0, triton 3.8.0. One decode step's admission at top-10 over 512 experts, into a 2,100-slot cache that already holds all 512, so every step is all hits and the time is the admission alone. 48 admissions, one per MoE layer, are captured in a CUDA graph and replayed 40 times; main's call and load_experts alternate A B B A. Microseconds per admission:

requests	routes	main	this PR
1	10	7.8, 9.8	8.6, 9.1
4	40	10.3, 7.5	6.9, 7.0
8	80	16.9, 16.7	14.7, 15.6
12	120	10.6, 11.1	11.7, 11.2
13	130	44.7, 43.2	45.3, 44.4
16	160	43.6, 44.2	44.5, 44.1
24	240	45.4, 44.9	44.9, 44.8
32	320	264.9, 264.4	75.1, 73.9
48	480	264.6, 263.8	93.2, 90.3
128	1,280	does not compile	92.2, 92.4

Splitting every step past 128 routes instead cost 59-70 µs from 13 to 24 requests, more than the one call it replaced, which is why one call keeps up to 256. Compiling main's 512-wide variant took 45 s on its first call.

Tests

• tests/moe/test_offload.py::test_offload_moe_layer_wide_decode_step_loads_its_experts_128_at_a_time: a 150-row step routing 300 times to 160 experts, 7 of them resident, admits them in queries of 128 and 32, each followed by its copy, and decodes at their slots. The statistics read 160 met, 153 missed and one call, and the histogram counts each route once. An 8-route step stays one admission, in place.
• test_dsv4_offload_layer_wide_on_demand_chunk_loads_its_experts_128_at_a_time: DeepSeek-V4's on-demand chunk, 300 routes over 640 experts, admits in queries of 128, 128 and 44.
• tests/moe/test_nvfp4_backends.py::test_triton_wide_decode_step_matches_the_streamed_layer_bit_for_bit (GPU): a 600-row step, 1,200 routes, over a 1,100-slot cache and over a cache of one layer's experts filled three per query. The GEMV over the slots equals the one over the streamed layer exactly, and the statistics count one step.
• test_triton_wide_decode_step_replays_as_a_cuda_graph (GPU): with both widths at 3, the wide admission and the GEMV capture as a CUDA graph, and a replay on a cold cache decodes exactly as the streamed layer.

Over main the two layer tests fail on a single query of 300 (assert [300] == [128, 32], assert [300] == [128, 128, 44]) and the GPU tests on the missing method. Through main's admission, the 600-row step fails to compile _lru_ensure_kernel with the error above.

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, tests/engine/test_cache_budget.py, tests/models/deepseek_v41 and tests/scheduler/test_scheduler_kv_usage.py.

• Linux, WSL2 Ubuntu 24.04 (Python 3.12, torch 2.14.0, triton 3.8.0), over main at 9c635a1: 166 passed, 7 skipped, the three tests that need flashinfer, which this environment lacks and which fail on main too (162 passed, 7 skipped), and one 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.
• 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): 156 passed, 6 skipped, 16 tests/models/deepseek_v41 tests that fail the same way on main, and one of the batch-memcpy tests.

The tracker has no issue or PR on a wide decode step's admission: searched lru_ensure, numel exceeds triton maximum, max-running-requests, ensure_experts batch, slot cache.

Known limits

• The measurement times the admission alone, every step all hits. The copies and the GEMV are the same either way and are left out.
• 256 is where one call stops being cheaper on this card; a card with cheaper launches may want to split earlier.
• The bit-identity tests run the Triton NVFP4 decode GEMV. Marlin and b12x read the same slot ids; neither runs here, since Marlin needs vLLM and b12x an sm_120 card.
• Hybrid decode is unchanged: its admission kernel loops over the query instead of deduplicating it in a block.