[jason-fxz]

Summary

Speed up victim selection in _ensure_experts_hybrid_kernel without changing its results, fix a race in its route rewrite, and add FREETOKEN_HYBRID_FETCH_ROUND to choose how many misses a hybrid decode step fetches.

Why main spills

Phase 2 loads the whole slot pool (next_pow2(pool slots) wide) into the registers of one CTA, then scans it once per route to mask the slots this step uses. On a 23,197-slot pool (Qwen3.8-Flash-Next NVFP4, 80 GB) that is 32,768 int64 usages plus ids: the cubin uses 32 registers and a 2,472 B stack, and 63% of its SASS is STL/LDL. One fetch costs 77 us inside the kernel.

Change

• No per-route mask: Phase 1 already sets the hit slots' usage to the current step, the largest usage in the pool, so after a CTA barrier the argmin never picks them.
• A pool of up to 8192 slots stays in registers; a larger one is scanned from memory in 8192-wide tiles per victim. Ties still go to the lowest slot.
• The victim's old owner is one scalar load; 8 warps for 4096-wide or wider tiles (with 4 warps the 8192 tile spills on sm80-sm89 when cross-compiled).

Routes, copy plan and cache state are bit-identical to main (lockstep over the same random routes, pools of 2,048 to 23,197 slots and per-layer pools).

Fix

Phase 3 rewrites expert_ids[i] in place. In the PTX every thread loads expert_ids[i] and then indexes slot_for_id with it, while only thread 0 stores the rewritten slot id; a warp that loads after that store indexes with a slot id. A barrier between the loads and the store fixes it. Without it, a 23,197-slot pool with clamp=False and back-to-back calls (no copy_missing in between) faults with an illegal memory access on main; with only this barrier added to main's kernel the same 72,000 calls run.

FREETOKEN_HYBRID_FETCH_ROUND

For M misses and the benched fetch fraction f (PCIe / CPU bandwidth), the step fetches about f * M experts and the CPU computes the rest:

• min-wait (default, current behaviour): of floor(f * M) and floor(f * M) + 1, take the one that makes the PCIe fetch and the CPU compute finish closest together.
• stochastic: floor(f * M), plus one with probability equal to the fractional part.
• prefer-fetch: ceil(f * M), so any miss fetches at least one expert into the cache.

Test

test_hybrid_ensure_matches_lru_reference runs the kernel against a Python model of the LRU contract, call by call from an empty to a full pool, for a pool that fits one tile, one that needs two, and per-layer pools, with clamp on and off and all three rounding modes. Before the barrier, the full parametrized run faulted in the two-tile clamp=False cases.

Results

H100 80GB, Xeon 8480+, CUDA 13.0, torch 2.14.1, Triton 3.8.0. nvidia/Qwen3.8-Flash-Next-NVFP4, AIME25 problem 1, T 1.0 / top-p 0.95 / top-k 20, decode tok/s median of 5 runs after one warm-up. A/B on main cc56b6b; after rebasing onto c5b5b72 the lockstep and the tests were run again.

slots	main	this PR	stochastic	prefer-fetch
23,197	100.2	100.1	104.7	106.1
18,432	100.2	101.8	103.9	101.7

Kernel time for one fetch at 23,197 slots: 76.9 -> 7.7 us. With the default rule a single miss is not fetched, so at 23,197 slots the kernel rarely reaches Phase 2.

ft bench bw --model nvidia/Qwen3.8-Flash-Next-NVFP4 -o benchbw.json
FREETOKEN_BENCHBW_PATH=benchbw.json [FREETOKEN_HYBRID_FETCH_ROUND=stochastic|prefer-fetch] \
python -m freetoken --model-path nvidia/Qwen3.8-Flash-Next-NVFP4 --port 2341 --served-model-name default \
  --max-running-requests 1 --cuda-graph-max-bs 1 --max-seq-len-override 40960 --kv-reserve-tokens 40960 \
  --cache-type naive --moe-strategy hybrid --moe-cache-size 23197 --expert-load parallel