Skip to content

fix(gemma4): the serving streams wait for what setup enqueued - #1030

Merged
FeathBow merged 1 commit into
pegainfer-project:mainfrom
FeathBow:fix/gemma4-load-handoff-ordering
Sep 2, 2026
Merged

fix(gemma4): the serving streams wait for what setup enqueued#1030
FeathBow merged 1 commit into
pegainfer-project:mainfrom
FeathBow:fix/gemma4-load-handoff-ordering

Conversation

@FeathBow

@FeathBow FeathBow commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator

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_safetensors synchronizes its stream after the unmap, so the host cost still overlaps the kernels and only the residual is paid; the log line and the LoadStats doc now say the expert kernels are drained, and elapsed_ms includes that residual.

The engine synchronizes before the lane exists. The ctx.sync() that sat inside the graph branch runs unconditionally before AsyncPrefillLane::new, so the KV pools, the arena and the warm passes are complete for every stream that follows, graphs or not.

Test Env

  • Single GPU (sm_89, x86_64), --release --features gemma4, the pinned Gemma 4 12B checkpoint and the 26B-A4B NVFP4 checkpoint.

Verification

  • cargo fmt --check clean; clippy --all-targets -D warnings clean for pegainfer-gemma4; lib tests 49 passed / 23 ignored, unchanged from main.
  • Lifecycle gates through the maintainer runner: 5 of 5 (the shared, green and gathered lane lifecycles on the dense profile, the shared and green ones again on the routed profile, which is the checkpoint with experts).
  • Same-config serving A/B on the 26B-A4B checkpoint, main vs this branch alternated twice each: four greedy prompts byte-identical across all eight runs (one md5 over text, finish reason and usage), with the default CUDA graphs and with --cuda-graph=false plus PEGAINFER_ASYNC_PREFILL=shared, the latter being the combination that had no ordering before.
  • Load time on the routed checkpoint (submission total from the loader's log line, two runs each, warm page cache): graphs on, main 10594 ms [10579, 10609] vs branch 10702 ms [10602, 10801]; graphs off with the lane, main 10645 ms [10641, 10648] vs branch 10701 ms [10695, 10706]. The difference is the residual expert-kernel drain and sits inside the run-to-run spread; the unmap's host cost still covers most of the kernels.

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>
@FeathBow
FeathBow merged commit ea02a9f into pegainfer-project:main Sep 2, 2026
17 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

gemma4: the serving streams do not wait for what the loader and setup enqueued

1 participant