[riony]

Part of #671. The prefill staging geometry this covers came in #689, whose review noted that "the new pool tests are all CUDA-only" and that "all current geometry coverage is CUDA-only". This test runs that geometry on CPU, where those tests skip.

What it is

Test-only: adds test_prefill_staging_unregisters_exactly_the_slots_it_overwrites to tests/moe/test_offload.py (72 lines, extending the existing file). No runtime change.

What it checks

• Builds a real OffloadMoeCache on CPU with slot pools of mixed row widths, including a pool whose MMQ kernel tail reads past its last row.
• Every pool's kernel tail bytes stay clear of the next pool's first row and of the bank end.
• Populates the slots the way decode would leave them (the decode kernel is CUDA-only), then runs the overlapped prefill.
• After that prefill: the staging overwrites exactly the first n slots of each pool, _pool_stable_from equals n, and only untouched slots stay registered, including a slot the staging only partly overwrites.
• Expectations come from the slot bytes before and after, not from the cache's own slot arithmetic.

Why CPU

All four slot-pool tests from #689 are @pytest.mark.skipif(not torch.cuda.is_available()) and build their caches on torch.device("cuda"), so a host without CUDA runs none of them, and no CPU-reachable test covers the staging slot count or _pool_stable_from. The review of #689 suggested extracting the arithmetic into a pure function to allow CPU unit tests; this drives the integrated path instead, which also covers the registration bookkeeping and a straddled slot.

Tested on

• GPU: NVIDIA RTX 3090 Ti, 24 GB; the host also has 2x RTX A4000, not used
• CPU: 2x AMD EPYC 7642
• RAM: 472 GiB
• OS: Ubuntu 26.04, kernel 7.0.0-43, NVIDIA driver 595.99.02
• FreeToken: upstream main at fff1022a, plus this branch

Commands, run from the repo root of each tree (base, then this branch), extensions built in-tree first with python setup.py build_ext --inplace:

# CUDA hidden
CUDA_VISIBLE_DEVICES= PYTHONPATH=python \
  numactl --cpunodebind=0 --physcpubind=0-15 --membind=0 \
  python -B -m pytest -p no:cacheprovider -rA -v --tb=long tests/moe/test_offload.py

# CUDA visible (RTX 3090 Ti only)
CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=GPU-15a57e09-1a95-37d7-679b-955eb2bf9079 PYTHONPATH=python \
  numactl --cpunodebind=0 --physcpubind=0-15 --membind=0 \
  python -B -m pytest -p no:cacheprovider -rA -v --tb=long tests/moe/test_offload.py

The mutation runs use the CUDA-hidden command on this branch with one line of python/freetoken/moe/offload_cache.py changed.

Results

suite	base main	this branch
tests/moe/test_offload.py, CUDA hidden	22 passed, 30 skipped	23 passed, 30 skipped
tests/moe/test_offload.py, CUDA visible (RTX 3090 Ti)	52 passed	53 passed

Mutation check: each edit below was applied to this branch, and only the new test failed, at the assertion named:

edit	fails at	what it shows
pool offsets computed without the tail padding	tests/moe/test_offload.py:310	assert (576 + 20) <= 576: a 20-byte MMQ tail runs into the next pool
_pool_stable_from rounded down	:343	assert 7 == 8: the floor says 7, but 8 slots were overwritten
invalidation floor rounded down	:345	assert [7] == []: a partly overwritten slot stays registered

The staging buffers sit at the start of every bank, and hit-D2D trusts every slot from _pool_stable_from onward. If the slot count or the floor rounds down, the cache keeps a slot whose bytes the staging already overwrote, and the next hit on that slot serves the wrong expert rows with no error anywhere. The test compares each slot's bytes before and after the prefill, so the expected boundary comes from the bytes rather than the cache's own arithmetic, and a slot the staging only partly overwrites still has to come out unregistered.

AI assistance

AI assistance was used in drafting the test and this description. I ran and reviewed the change on the hardware above and can explain and defend every line of it without help.

This is my first contribution to FreeToken. I have also posted in the Developer Slack about a multi-GPU expert tier (related to #447) and would welcome a pointer on where that should integrate.