perf(gemma4): the routed experts load through pinned staging - #1023
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5931b16741
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| shards: &[SafeTensors], | ||
| manifest: &Manifest, | ||
| ) -> Result<Vec<Option<StackedExperts>>> { | ||
| let mut stager = ByteWeightStager::new(ctx)?; |
There was a problem hiding this comment.
For dense Gemma 4 checkpoints such as 12B and 31B, every layer.moe is None, but this eagerly constructs a ByteWeightStager anyway. Its constructor allocates two 32 MiB pinned buffers and creates a Rayon thread pool, so dense startup now pays an unnecessary 64 MiB pinned-memory allocation and can fail on hosts where the existing loader's pinned buffers already approach the available pinning limit, even though no expert upload will occur. Construct the byte stager only after determining that the manifest contains routed layers.
Useful? React with 👍 / 👎.
Signed-off-by: Feathbow <feathbow@gmail.com>
5931b16 to
d92398e
Compare
1 participant
Description
Closes #1022
The expert tensors stop paying one pageable copy per slice. The A4B expert packed weights and scales used to reach the device as one pageable
htod_copyper source slice — 26,677 serialized submissions on a 26B load, 22.2 s of API time warm and 45.1 s cold (98.5% of total API time) on a profile of the startup span, while the copy engine needs well under a second for the bytes.They ride the loader's existing pinned double buffers instead. A
ByteWeightStagergives the stager a raw-byte destination; per-expert slices coalesce into staged chunks (next_contiguous_chunk, unit-tested) and upload through the pinned pool with the usual fence-and-swap. The e4m3 scale peak folds online during staging, so the whole-projectionscale_byteshost copy goes away. The GPU repack is untouched, the per-expert length assertions stay, and the API is additive — other model lines are unaffected.The claim is the submission collapse, not the warm wall. Measured on the same span: 1,421 submissions, 11 ms API time, 709 ms of copy-engine time (~31 GB/s, link speed). Warm ready-to-serve wall is parity — main 12.2 s [12.2, 12.2] vs this branch 12.2 s [11.1, 13.2] over three alternating rounds — because the transfer proper is a small share of a warm ready; the serialization is what dominates the cold path.
Test Env
--release --features gemma4, Gemma-4-26B-A4B NVFP4 checkpoint.Verification
ignore_eos) produce byte-identical outputs on main and on this branch (matching md5 over text, finish reason and usage).cargo fmt --checkandclippy -D warningsgreen; the staging chunker carries a host unit test (contiguous_chunks_preserve_source_bytes).