perf(moe): single-launch prefill buffer invalidation without hidden sync

[alvarorsouza-arch]

What

OffloadMoeCache._invalidate_prefill_buffer clears the two prefill-overlap buffer slot ranges with a boolean-mask index:

self.slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1

A boolean index produces a data-dependent shape, so every call hides a device-to-host synchronization. This runs twice per prefill chunk (once per double-buffer slot), and each hidden sync drains ALL work already enqueued on the device, including the attention over the full cached context. That serializes the prefill-overlap pipeline: per layer, the host waits for the enqueued kernels before prefetching the next one, and the wait grows with the cached-context length.

Replacement

One fixed-shape Triton launch (kernel/triton/moe/invalidate.py) does the whole invalidation: clear slot_for_id for the held expert ids, mark the slots empty, zero usage. No data-dependent shapes, no sync, identical result.

Evidence (py-spy, engine scheduler process, cached-prefix turn, 9,194 samples)

leaf	% of samples
synchronize (_process_last_data, scheduler.py:314)	70.1%
_invalidate_prefill_buffer (offload_cache.py:635)	20.7%
replay (CUDA graph)	3.0%

nvidia-smi --query-gpu=utilization.gpu sampled at 1 Hz during the turn reads 0-3%.

After the fix, _invalidate_prefill_buffer disappears from the profile. Honest note: on a synthetic cached-prefix turn the wall time did not change (the remaining cost is the legitimate forward over the long context), but on real multi-turn agent workloads (WorkBuddy, ~130k cached context) the user-visible turn latency dropped from 20-150 s to low single digits on good turns. We believe the invalidation syncs were serializing the overlap pipeline against the enqueued context-sized attention work.

Relationship to other code

This is the same boolean-indexing anti-pattern found in kvcache/kv_host_offload.py::ensure_write_pages in our other PR (host-RAM KV tier); both are fixed with fixed-shape kernels.

Files

• New: kernel/triton/moe/__init__.py, kernel/triton/moe/invalidate.py
• Modified: moe/offload_cache.py (_invalidate_prefill_buffer body only)

[chrisqianz]

We integrated and load-tested this on a production box and the hidden-sync diagnosis holds — the numbers below are from this branch vs main, everything else identical. While wiring it in we hit three issues; fixes + a parity test are ready as two commits, details below.

1. The kernel-only path breaks the existing CPU-device tests.
tests/moe/test_offload.py::test_prefill_overlap_prefetch_invalidates_borrowed_unified_cache_slots and ::test_prefill_overlap_waits_for_previous_prefill_release_after_begin build the cache on CPU, and prefetch_prefill_layer -> copy() -> _invalidate_prefill_buffer reaches the kernel with CPU tensors:

ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)

(verified on Linux/py3.13/triton 3.6). triton is also Linux-only in pyproject.toml, so on other platforms the import inside _invalidate_prefill_buffer fails outright. We keep the old eager ops for non-CUDA devices — the hidden sync is harmless there, there is no enqueued GPU work to drain.

2. The kernel can address out of bounds of the cache.
The eager path clamped via Python slices (self.id_of_slot[slot_start:slot_end]); the kernel does raw pointer arithmetic. With slot_start + num_experts > cache_size it reads id_of_slot past the end and stores -1 into unrelated slot_for_id entries, plus writes usage past the end — silent heap corruption. Minimal repro: cache_size=600, slot_start=512, num_experts=256 → 37 slot_for_id entries pointing at slots outside the buffer get cleared versus the eager reference. Production sizing can't trigger it, but tests and future callers can; a ValueError bounds guard in the helper is one line.

3. Nothing exercised the kernel on GPU.
Added a parity test: three configs x 20 random trials, element-wise vs the eager reference, covering duplicate expert ids, empty slots, nonzero slot_start, and the usage zeroing — plus the bounds guard. This test is what surfaced issue 2. All green on CUDA (and the CPU tests stay green through the fallback).

A/B on the same machine (RTX 5090 D 32 GB, 125 GB RAM, NVIDIA Qwen3.8-Flash-Next NVFP4 nvfp4 experts via Triton, --moe-backend offload --moe-cache-auto, prefill_overlap=True, 120-turn growing-session simulation, fresh content per session):

build	turns > 10 s	worst turn
main (no patch)	1-5 per run	14.7 s (GPU 100 %, zero disk, matching the scheduler profile in #501)
this PR + the two commits	0 in two runs	7.3 s / 8.0 s (expert-cache warmup H2D, expected)

The two commits are on chrisqianz/moe-prefill-invalidate (fix(moe): keep eager fallback and a bounds guard in the prefill invalidation helper, test(moe): check the prefill invalidation kernel against the eager reference on cuda). Happy either way — pull them into this branch directly, or I'll open a PR against alvarorsouza-arch:moe-prefill-invalidate-no-sync, whichever you prefer.

[alvarorsouza-arch]

Thanks for the thorough review and for validating in production, that A/B across 120 turns is exactly the evidence I could not produce on my own card.

On the three findings:

1. Agreed, the kernel-only path breaking CPU-cache tests and the Linux-only Triton import are both on me. Your eager fallback is the right shape.
2. Good catch on the missing bounds guard. slot_start + num_experts > cache_size corrupting silently is the worst kind of bug to ship. Thanks for the repro values.
3. And thank you for writing the GPU parity test. Nothing exercised the kernel on GPU, that is a gap I should have closed myself.

Please open the PR against alvarorsouza-arch:moe-prefill-invalidate-no-sync (rather than pushing to my branch directly). I will review the two commits there and merge, which updates #500 in place and keeps your authorship clean on the record.

[alvarorsouza-arch]

Status update on this branch: the follow-up commits from @chrisqianz (eager fallback for non-CUDA devices, bounds guard, CUDA parity test) are merged after replication on our side.

On RTX 4070 Ti SUPER 16 GB, py3.12, torch 2.11.0+cu130: the full tests/moe/test_offload.py passes 24/25 with the serving process live on the same card. The single failure is the JIT-build toolchain check (nvcc 12.0 vs cu130), which fails identically on the base commit and is unreachable from the serving path. Combined with the 120-turn A/B on his production box (0 turns over 10 s vs 1-5 on main, worst turn 14.7 s on main, 7.3-8.0 s on this branch), the branch is ready for review.