perf: fix overlap scheduling and all-reduce fusion for NVIDIA Confidential Computing(CC) on Blackwell - #31447
Open
elvischenv wants to merge 2 commits into
Open
perf: fix overlap scheduling and all-reduce fusion for NVIDIA Confidential Computing(CC) on Blackwell#31447elvischenv wants to merge 2 commits into
elvischenv wants to merge 2 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces support for NVIDIA Confidential Computing (CC) in SGLang. It implements an AsyncD2HCopyWorker to offload device-to-host copies to a dedicated background thread, preventing scheduler thread stalls caused by synchronous bounce-buffer copies under CC. It also configures FlashInfer allreduce fusion to force the multicast-free trtllm backend and disables symmetric-memory preflight checks when CC is enabled. The review feedback suggests unconditionally initializing enable_async_d2h_copy and async_d2h_worker in init_overlap to prevent potential AttributeErrors when overlap scheduling is disabled.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
elvischenv
changed the title
elvischenv
force-pushed
the
cc-fixes-rebased
branch
from
July 17, 2026 06:44
420cc77 to
cb22a38
Compare
elvischenv
force-pushed
the
cc-fixes-rebased
branch
from
July 28, 2026 03:02
cb22a38 to
befa9d9
Compare
elvischenv
marked this pull request as ready for review
July 28, 2026 03:03
elvischenv
requested review from
Ying1123,
hnyls2002,
merrymercy and
xiezhq-hermann
as code owners
July 28, 2026 03:03
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
elvischenv
requested review from
BBuf,
Edwardf0t1,
Fridge003,
HaiShaw,
ch-wan and
ispobock
as code owners
July 28, 2026 03:03
Collaborator
|
/tag-and-rerun-ci |
elvischenv
force-pushed
the
cc-fixes-rebased
branch
from
August 13, 2026 03:43
befa9d9 to
383e2cc
Compare
…er CC Under NVIDIA Confidential Computing (bounce-buffer CC) a device-to-host cudaMemcpyAsync is forced synchronous and blocks at issue. The overlap scheduler's per-step result readback (next_token_ids, logprobs, hidden states, ...) then stalls the scheduler thread for ~a full copy every step, serializing decode and collapsing overlap. Detect CC once via NVML (is_confidential_compute, env-overridable with SGLANG_CONFIDENTIAL_COMPUTE) and, when overlap is enabled, offload the blocking readback to a dedicated daemon thread (AsyncD2HCopyWorker) that owns a private CUDA stream. Mirrors the TensorRT-LLM pattern (NVIDIA/TensorRT-LLM#8463): - submit() records a readiness event on the caller's current stream and hands the copy to the worker; the worker event-syncs on it (event-sync, not a stream wait, so the blocking copy never stalls the scheduler's CUDA API calls), runs the copy on its private stream, and signals it; - submit() returns a HostCopyDone -- a host-thread-backed drop-in for the copy_done CUDA event (same record()/synchronize() surface) -- which the scheduler stores in result.copy_done. Every existing copy_done.synchronize() consumer then waits on the worker unchanged, so no result-processing code changes; - synchronize() re-raises if the copy failed, so a readback error aborts the step instead of consuming invalid host tensors. The stream is created and owned by the worker so nothing else can enqueue onto it (sharing it would recouple copy-completion to the next forward and defeat the overlap). Only GenerationBatchResult is offloaded -- its copy_done is pre-created, so the swap survives copy_to_cpu; embedding is one-shot with no decode overlap to preserve and is copied inline. Co-authored-by: spethe <spethe@nvidia.com> Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ting
The trtllm AR+RMSNorm fusion runs off-CC, but under CC it disabled itself:
create_allreduce_fusion_workspace allocated a symmetric-memory (cuMulticast)
workspace whose preflight fails under CC. The fusion kernels are themselves
multicast-free (one-shot Lamport and two-shot sync; 0 multimem in
trtllm_allreduce_fusion.cuh) -- only the workspace allocator wanted multicast.
FlashInfer now auto-selects a multicast-free IPC workspace under CC, so sglang
reuses the normal off-CC fusion path:
- _resolve_backend forces the trtllm backend under CC (the SM100 "auto"
default is mnnvl, which needs NVLink multicast and is unavailable under CC)
and raises on an explicit mnnvl request or multi-node;
- skip the symmetric-memory preflight under CC (its fabric / cuMemCreate
probe is exactly what CC blocks);
- the workspace stays a normal AllReduceFusionWorkspace object, so
is_buffer_size_sufficient / cleanup / allreduce_fusion use the standard
object path -- no cc-specific IPC field, helper, or branch.
Enabled by default whenever CC is detected (is_confidential_compute()); no env
var needed, and off-CC behavior is unchanged. Requires the matching FlashInfer
CC auto-detection change (validated CC-on for both the one-shot and two-shot
kernels).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
elvischenv
force-pushed
the
cc-fixes-rebased
branch
from
August 14, 2026 02:26
383e2cc to
64121c6
Compare
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
2 participants
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.
Motivation
Two fixes so SGLang runs efficiently under NVIDIA Confidential Computing (CC), where the GPU operates behind a bounce buffer:
cudaMemcpyAsync(D2H) is forced synchronous — it blocks at issue — so the per-step readback ofnext_token_ids/ logprobs serializes onto the scheduler's critical path and kills decode overlap.create_allreduce_fusion_workspaceallocates a symmetric-memory (cuMulticast) workspace whose preflight fails under CC, so the fusion silently falls back off — even though the fusion kernels themselves are multicast-free.Both paths are gated on
is_confidential_compute()(NVML, overridable viaSGLANG_CONFIDENTIAL_COMPUTE); off-CC behavior is unchanged.Modifications
cc: async device->host result-readback workerpython/sglang/srt/managers/async_d2h_copy_worker.py(AsyncD2HCopyWorker): runs the per-step D2H copy on a dedicated thread with its own CUDA stream, off the scheduler's critical path, so overlap is preserved. Wired intoscheduler.py/managers/utils.py; only routed through the worker when CC is detected.is_confidential_compute()tosrt/utils/common.py.cc(fix): enable FlashInfer AR+RMSNorm fusion under Confidential Computing(layers/flashinfer_comm_fusion.py)_resolve_backendforces thetrtllmbackend under CC (SM100 "auto" defaults tomnnvl, which needs NVLink multicast), the symmetric-memory preflight is skipped, and the workspace stays a normalAllReduceFusionWorkspace. Enabled by default whenever CC is detected.Accuracy Tests
With CC on:
Checklist
test/registered/core/test_async_d2h_copy_worker.py).🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #31763733100
Latest PR Test (Extra): ❌ Run #31763732864