feat(moe): split the offload expert cache into slot pools

[jason-fxz]

Part of #671.

Summary

The offload expert cache used one LRU pool for all MoE layers. This PR splits it into slot pools: a pool is a contiguous slot range with its own LRU, each MoE layer belongs to one pool, and a pool holds layers of one quant.

Arg	Values	Meaning
--moe-cache-scope	unified (default), per-layer	One pool per quant, or one pool per layer. On a single-quant model unified is one pool, as on main.
--moe-cache-split	equal (default), profile:<path>	Slots per pool by layer count, or by expert use from a profile.
--moe-cache-mem	e.g. 12GiB	Cache size as memory instead of a slot count.

The cache also holds pools of different quants, e.g. GGUF Q4_K_M with ffn_down in Q6_K on some layers and Q4_K on others. No loader produces such banks yet.

Changes

Cache (moe/offload_cache.py)

• Pools: pool_of_layer / pool_sizes, pool-local slot ids, per-pool LRU state and bank views.
• Several quants: a layout key per layer (quant tag + bank dtype and row shape); typed banks for one layout, byte segments for several.
• Prefill staging: buffers at the start of every bank, a static invalidation set per buffer, a hit-D2D threshold per pool.
• With pools of several quants and a hybrid kernel that reads inactive routes, staged slot 0 rows are zeroed after prefill.

Kernels (moe/offload_kernels.py)

• Expert LRU, hybrid ensure and prefill hit lookup run on the layer's pool.
• The hybrid fetch count is capped at the pool's free slots.

Policy (moe/pool_policy.py, new)

• group_layers (scope), make_pools, split_slots (split), profile loading.

Engine and args

• New args; ExpertBanks.layer_tags for loaders.
• Offload with a pool below num_experts slots lowers max_running_req and the CUDA graph batch sizes to fit one decode step.
• Several pools do not combine with --moe-strategy cpu or --moe-cpu-layers.
• A runtime rebuild re-splits the pools.

Other

• DSV4: the small-chunk prefill threshold uses the smallest pool.
• ft bench bw gather rig and cache status follow the new bank storage.

Validation

• Unit tests (H100): tests/moe tests/engine tests/server tests/scheduler tests/layers tests/dsv4 tests/kvcache 1271 passed, 12 skipped. New tests cover per-layer pools, pools of two layouts through prefill, quant tags, rebuild memory, the split policy and the args.
• E2E on earlier revisions (before f02fbb1), greedy:
  • Qwen3-30B-A3B bf16, unified, offload and hybrid: tokens identical to main.
  • Qwen3.6-35B-A3B bf16 offload, six prompts: per-layer equal vs unified miss rate 42.79% vs 44.33% at 1024 slots, 48.70% vs 51.12% at 768 slots, same tokens. The profile split was 0.3-0.5 pp worse than equal.
  • DeepSeek-V4-Flash-0731 offload: per-layer tokens identical to unified.
• Not run yet: E2E on f02fbb1, --moe-cache-mem end to end, offload with more than one request.

[andy-yang-1]

Static review only; I didn't run it locally (the new pool tests are all CUDA-only). I didn't find anything that produces wrong results, and the single-pool path looks equivalent to main. A few things I'd like addressed or discussed before merging, roughly in priority order.

1. The multi-layout (per-quant) path is unreachable from the engine

Every format that binds a quant method (bf16, fp8_block, nvfp4*, mxfp4) passes a single method.layout(), and set_bank_sources checks every layer against it:

FreeToken/python/freetoken/moe/offload_cache.py

Lines 486 to 492 in 15f3bff

	if self.layout is not None:
	spec = self.layout[name]
	if tuple(source.shape[1:]) != tuple(spec.shape) or source.dtype != spec.dtype:
	raise ValueError(
	f"bank {name!r} rows are {tuple(source.shape[1:])} {source.dtype} but the expert kernel's layout "
	f"wants {tuple(spec.shape)} {spec.dtype}; the banks were packed for another kernel"
	)

On top of that, shared_offload_method raises when layers disagree on (kind, kernel). So the raw byte-segment mode can only show up for GGUF (no method, layout=None), and _expert_gemm only dispatches GGUF on quant_format == "q4_0" today.

That means raw banks, _stage_slot0_rows, zero_inactive_route_slots, the per-layout gather pointers and layer_tags (a sizable share of the cache changes) are exercised only by tests. When a Q4_K + Q6_K loader lands, the layout check will need to become per tag / per pool and the kernel dispatch per layer anyway.

Suggestion: move the multi-layout support into the follow-up that adds a real loader, or at least state this limitation in the code.

2. max_running_req gets silently clamped very low

The startup split only reserves room for one token per step:

FreeToken/python/freetoken/engine/engine.py

Lines 763 to 764 in 15f3bff

	# with GPU decode, a pool below num_experts slots bounds the batch; the startup split keeps room for one token per step
	pools = self._moe_pools(1)

Then, if any pool is below num_experts, _clamp_decode_batch_to_pools lowers max_running_req and the CUDA graph batch sizes to min_pool // top_k, with only a warning.

• per-layer + equal: e.g. Qwen3-30B-A3B (48 layers, 128 experts, top-8) at 1024 slots gives ~21 slots per pool, so concurrency drops to 2.
• profile split: the coldest layer almost always ends up with exactly top_k slots, so concurrency drops to 1.
• The PR description lists "offload with more than one request" as not run yet.

Suggestion: for GPU decode, split with _moe_pools(config.max_running_req) so the configured concurrency is honored, and fail (or require an explicit opt-in) when the budget can't cover it, rather than degrading silently.

3. The hybrid ensure kernel recompiles per pool size

_ensure_experts_hybrid_kernel declares cache_size: tl.constexpr and now receives each pool's slot count:

FreeToken/python/freetoken/moe/offload_kernels.py

Lines 293 to 294 in 15f3bff

	num_experts: tl.constexpr,
	cache_size: tl.constexpr,

Each distinct pool size is a separate Triton JIT compile. Equal split yields at most two sizes, but a profile split can yield up to one per layer, which noticeably slows the first decode and graph capture. cache_size is only used for the mask and the new fetch cap, so making it a runtime argument (do_not_specialize) while keeping BLOCK_C constexpr should be enough. flashlib's lru_ensure already does this for num_cached.

4. --moe-cache-mem can hit an assertion

FreeToken/python/freetoken/engine/engine.py

Lines 783 to 789 in 15f3bff

	elif config.moe_cache_mem is not None:
	per_slot = expert_bytes_per_slot(banks.sources)
	size = min(config.moe_cache_mem // per_slot, size_cap)
	# per_slot is a mean over layers, so with pools of different slot sizes the split can exceed the budget: shrink until it fits
	while size > 0 and self._moe_pool_bytes(split_slots(size, pools, self._moe_expert_counts)) > config.moe_cache_mem:
	size -= 1
	object.__setattr__(config, "moe_cache_size", size)

If the resolved slot count lands in [E, 2E) with prefill overlap left at its default (on), the cache_size >= 2 * num_experts assert in OffloadMoeCache.__post_init__ fires. The --moe-cache-auto path turns overlap off in that case; the mem path doesn't. An explicit --moe-cache-size behaves the same way on main, but mem users can't predict the slot count. Suggestion: mirror the auto path, i.e. disable overlap and log it.

5. The profile path is validated late

resolve_profile runs after load_expert_banks, so a typo in profile:<path> only fails after all expert banks have loaded, which can take minutes. Checking that the file exists in args.py would catch it up front.

Nits

• _slice_pools is hard to read, e.g. the _pool_stable_from one-liner below. It's pure arithmetic; pulling it into a pure function (with a _ceil_div helper) would also allow CPU unit tests, since all current geometry coverage is CUDA-only.
  

  FreeToken/python/freetoken/moe/offload_cache.py

  Lines 369 to 372 in 15f3bff

  	self._pool_stable_from = [
  	max(min(max(-(-(2 * stage[k] - offsets[k][p]) // rows[k][p]), 0), s) for k in range(len(rows)))
  	for p, s in enumerate(self.pool_sizes)
  	]

• layer_tags entries are used as dict keys, so they must be hashable. Document it or normalize them to tuples.
• Layers with the same byte layout but different quants must pass layer_tags; otherwise they share a pool and hybrid's slot-0 read may decode another quant's bytes into NaN. Nothing enforces this.
• decode_routing_stats still uses cache_size / num_layers, which doesn't reflect pools.
• With several pools, materialize_layer allocates torch.arange(E) per layer per prefill; it could be cached.
• A runtime rebuild never raises max_running_req back, and with GPU decode it requires every pool to hold min(E, top_k * max_running_req) slots, so per-layer caches can barely shrink at runtime. Worth documenting.

Checked and looks fine

• Single-pool behavior matches main: materialize, copy destinations, staging views, the hit-D2D threshold (min(2E, cache_size)), the new hybrid fetch cap (never binds with one pool), and the DSV4 small-chunk threshold (still E).
• Pool-local slot ids are used consistently in lru_ensure, the hybrid kernel, copy_missing's per-pool dst pointers, alphas_for_slots and the hit gather.
• The staging invalidation set is the union over banks and covers every slot below _pool_stable_from in each pool, so hit/miss classification is safe. Hand-computing the two-layout test case gives [6, 5, 8, 0], matching the assertion.
• Slot-0 zeroing only happens after those slots are invalidated, at the last layer's release on the compute stream, ordered after both staging copies.
• The clamp runs before the KV pool, page table and scheduler read max_running_req, and the scheduler shares the same config object.

[jason-fxz]

Thanks for the review, pushed 8fa6381, cb2acd8 and 400cb01 on top of 15f3bff.

1. Kept: the per-quant pools are for GGUF, whose experts can use a different ggml type per layer, and a GGUF branch stacked on this PR already runs on them.
2. Kept as is: it only applies to offload with a pool below num_experts, logs a warning and is documented under --moe-cache-split, while splitting with the default max_running_req of 4 would refuse per-layer pools on Qwen3.6-35B-A3B below 1280 slots.
3. Fixed in 8fa6381: cache_size, id_of_slot_ptr and usage_ptr are no longer specialized, so the 40 pools of a Qwen3.6-35B-A3B profile split compile 4 times (0.82 s) instead of 33 (7.76 s), with identical outputs and no slowdown in hybrid end-to-end runs.
4. Fixed in cb2acd8: below 2 * num_experts slots the mem path turns prefill overlap off and logs it, like --moe-cache-auto.
5. Fixed in 400cb01: the profile is read before the expert banks load, so a bad path fails before them.

Nits: left as is (decode_routing_stats has no caller, on main either).