fix(gemma4): the serving streams wait for what setup enqueued - #1030
Merged
FeathBow merged 1 commit intoSep 2, 2026
Merged
Conversation
Both device contexts create non-blocking streams with cudarc's event tracking disabled, so nothing orders work across them unless the code does. The loader enqueued the expert repack and scale kernels on its stream and returned without draining them; the engine's own synchronize ran only under CUDA graphs, and the async prefill lane's stream was created with no dependency on the engine stream. What hid both was time: the shard unmap's host cost after the enqueue, and the startup work before the first request. The loader now drains its stream after the unmap, so only the residual is paid and every later consumer sees finished experts. The engine synchronizes its stream before the lane exists, graphs or not, so the KV pools, the arena and the warm passes are complete for every stream that follows. Signed-off-by: Feathbow <feathbow@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fixes #1029
Both handoffs relied on time. The loader's and the engine's
DeviceContexts each own a non-blocking stream with cudarc's event tracking disabled, so nothing orders work across them unless the code does. The loader returned with the expert repack and scale kernels still enqueued on its stream, and the engine created the async prefill lane's stream without synchronizing its own unless CUDA graphs were on. The shard unmap's host cost and the startup work before the first request were what kept the readers behind the writers.The loader drains before handing over.
from_safetensorssynchronizes its stream after the unmap, so the host cost still overlaps the kernels and only the residual is paid; the log line and theLoadStatsdoc now say the expert kernels are drained, andelapsed_msincludes that residual.The engine synchronizes before the lane exists. The
ctx.sync()that sat inside the graph branch runs unconditionally beforeAsyncPrefillLane::new, so the KV pools, the arena and the warm passes are complete for every stream that follows, graphs or not.Test Env
--release --features gemma4, the pinned Gemma 4 12B checkpoint and the 26B-A4B NVFP4 checkpoint.Verification
cargo fmt --checkclean;clippy --all-targets -D warningsclean forpegainfer-gemma4; lib tests 49 passed / 23 ignored, unchanged from main.--cuda-graph=falseplusPEGAINFER_ASYNC_PREFILL=shared, the latter being the combination that had no ordering before.