[kalebbroo]

What

Follow-up to #914. That PR's 4090 A/B showed the device expert cache about 16% slower than plain auto-promotion over 32 decode steps: 268.7 s vs 231.4 s.

The cause: with no host runner, every miss is uploaded for its call, and admission then read the same expert again after the layer. That came to 1178 admissions, about 14.8 GB of second reads, each a synchronous staged upload from the volume.

Now, for an offload with no host runner:

• The misses that the admission rule takes are uploaded into the cache before the layer runs them (MoeExpertOffload.AdmitNow, the same fills-then-hot-swaps rule as AfterLayer).
• They are pinned with a second lease, which waits for their uploads, and then run from the cached copy.
• Result: one read per admitted miss.
• A key whose upload did not start stays a per-call miss.

Offloads with a host runner (Qwen, GLM and the others) keep the after-layer admission, so nothing changes for them.

MoeOffloadStats counts an admitted-before-run expert's rows as streamed, not resident, so the resident share still means "served without an upload".

Validation

• CPU lane: LLM.Tests Ernie5 + Moe 56 passed. The new test gives the cache room for every expert. The first forward's misses are all admitted before they run, and the second forward runs every row resident, with no further admissions or streamed rows. Outputs match the direct path.
• Engine.Tests: Ernie5 synthetic 34 passed, 1 skipped (real checkpoint absent).
• Not run yet: the real-checkpoint 4090 A/B against no cache. The pod run follows; its table goes in a comment.

[kalebbroo]

@codex review

[kalebbroo]

@claude Review PR

[claude]

Claude finished @kalebbroo's task in 1m 28s —— View job

---

Review of PR #916: admitting a host-less offload's misses before they run

I found no bugs that block this. The approach works. On CUDA, BeginUpload sends the tensors through CudaStreamingWeightCache. The projection for a Placement == Cpu miss on the host-less path therefore finds the cached device copy once _offAdmitLease has pinned it, so each admitted miss really is read only once.

What I checked

• Leases: _offAdmitLease.Dispose() is safe when the lease was never bound, because Owner?.Release does nothing then, and it is safe when the lease was already released. If AcquireResident/AwaitPending throws, the lease is left untouched, and the finally still releases _offLease.
• Eviction safety: Prefetch adds the current layer to protectedLayers and puts the admitted keys in requested. The layer's resident experts are still pinned by _offLease. So a fill or swap can only evict from other layers, the same as AfterLayer did.
• Partial uploads: if Prefetch stops early (MakeRoom fails or OutOfVramException), the keys it did not start never get an entry. They go to _offAdmitMissed and run as per-call misses, as the PR description says.
• Admission scores: the scores SelectAdmissions reads already include this step's access, because Plan's AcquireResident notes an access for absent keys too. So the swap gate (AdmitMinScore) behaves as it did after the layer.
• Hosted offloads: offloads with a host runner (Qwen, GLM and the others) do not change; they still use AfterLayer.

Suggestions (non-blocking)

1. The uploads do not overlap resident compute. MoeFeedForward.cs:629-632: AcquireResident waits for every admitted upload (up to MaxFillsPerLayer = 16) before step 2 queues any expert. The resident experts are already pinned by _offLease, so they could run first:

   • start the uploads (AdmitNow) where they start now;
   • queue the Placement == Gpu experts;
   • then bind _offAdmitLease and run the misses.

   That way the staged uploads overlap the resident experts' GEMMs. This is worth trying if the 4090 A/B still shows the cache behind auto-promotion.

2. The pin inflates the cache's stats and eviction scores. Each admitted miss is now counted as a _prefetches entry, then as a hit (_hits) and often an in-flight hit (_inFlightHits) in PinInto. It also gets a second NoteAccess for a single routed use. The code comment mentions the score effect, but two things follow:

   • ExpertCacheStats hit rates go up on the host-less path, so they cannot be compared with hosted runs or with perf(ernie5): device expert cache for the routed experts #914's numbers.
   • Freshly admitted experts score +1 above residents with the same real use, which favours keeping them over older residents when the cache is full.

   If that matters, an internal AcquireResident overload that skips NoteAccess and hit accounting would fix both.

3. The docs describe the old behaviour.

   • The remarks on MoeExpertOffload (MoeExpertOffload.cs:30-39) still say admission happens "after a layer runs" and that callers report with Record and AfterLayer. They should mention AdmitNow for Host is null.
   • The StreamedRows doc (MoeExpertOffload.cs:9) says "uploaded for that one call". Admitted-before-run rows now count there too, so that line should say so, or the decision in the PR description will only be recorded in the PR.

4. Minor style. System.Runtime.InteropServices.CollectionsMarshal is fully qualified twice in MoeExpertOffload.cs; a using would be cleaner.

5. Test coverage.

   • The new test uses HostExpertCache, where BeginUpload returns null. That means the "pin waits for a pending upload" path (AwaitPending) and a partial Prefetch (budget too small for every admitted miss, so _offAdmitMissed is non-empty) are never exercised.
   • A second case with a tight budget would cover the fallback where "a key whose upload did not start stays a per-call miss". It would check that the output still matches and that StreamedRows grows.

Not done

• I did not build the code or run the tests here; this review comes from reading the source.
• The real-checkpoint 4090 A/B is still pending, as the PR notes.

TL;DR: The change is correct and safe to merge once the pod numbers look good. The main follow-up is to run resident experts while the admitted uploads are in flight. Also worth doing: update the StreamedRows and admission docs, and note that the cache's own hit stats are now higher on this path.

[chatgpt-codex-connector]

Codex Review Summary: ✅ Completed

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-15T07:43:46.890544Z	596bc99	Manual request

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Chef's kiss.

Reviewed commit: 33757b6651

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[kalebbroo]

Real-checkpoint A/B, RTX 4090 (new pod, 124 GB RAM, network volume, code at 33757b6). Both cases ran resident with tokens 1..5, then 32 greedy decode steps, on the same pod.

	No cache (auto-promotion)	Device expert cache, admit before run
Load	138.5 s	125.9 s
Prefill 5 tok	160.0 s	157.6 s
Decode, 32 steps total	219.4 s	243.6 s (+11.0%)
Decode, mean of last 8 steps	2.82 s	3.67 s
Greedy tokens	identical	identical
Per-layer parity vs reference	—	43/43 layers at cosine >= 0.999

• Gap vs perf(ernie5): device expert cache for the routed experts #914: the slowdown went from +16% (perf(ernie5): device expert cache for the routed experts #914, double read) to +11%.
• Admissions: 1174, about the same as before, each now read once.
• Hit share rises from 22% at the first step to 53% by step 32.
• Prefill fills the whole budget (+227 on the first decode step, then 0 for steps 2–7). With one swap per layer per step and an admit score of 8, the cache then adapts slowly.
• Budget gap: the no-cache run's auto-promotion fills VRAM to its own smaller headroom (about 22 GB used), while the cache stops at 10.5 GB, a 15% / 3 GiB headroom. So the baseline holds more experts.
• Next experiment: a smaller headroom for the decode-only cache, plus faster swaps. The cache stays opt-in and test-only, so no default path changes.

🤖 Addressed by Claude Code

[kalebbroo]

Addressed in 596bc99:

• Docs (3): the MoeExpertOffload remarks now describe admitting before the run with AdmitNow and the second lease when there is no host runner. The StreamedRows doc now says admitted-before-run rows count there.
• Style (4): CollectionsMarshal comes from a using.

Not changed:

1. Overlap with resident compute: a mapped expert's admission upload is synchronous on the host (staged positional reads, each chunk waited on before the buffer is reused), so starting it earlier would not overlap it with the resident GEMMs. Overlap needs an asynchronous mapped path first. That is a follow-up, together with the budget and swap-rate tuning the A/B points to.
2. Stats and score inflation: agreed. Hit rates on the host-less path cannot be compared with hosted runs. The A/B reads MoeOffloadStats, where these rows count as streamed. An internal pin that skips NoteAccess and hit accounting goes with the tuning follow-up.
3. Pending-upload and partial-prefetch coverage: HostExpertCache cannot exercise an asynchronous pending upload. The CUDA A/B ran this path on the real checkpoint (1174 admissions, parity holding on all 43 layers). A partial Prefetch is rare here, because the selection only fills free room. A CUDA-side unit test would cover it and is noted for the follow-up.

🤖 Addressed by Claude Code

[kalebbroo]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. 🎉

Reviewed commit: 596bc99c37

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".