Repository navigation
[None][fix] Drain guided-decoder host functions before GIL-holding forward calls - #19455
Conversation
|
/bot run |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughThe change drains pending guided-decoder host callbacks outside CUDA graph capture. It applies this behavior to execution paths and routes ChangesGuided decoder host-function draining
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~20 minutes Change: Bug fix · Severity of issue fixed: Medium Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
PR_Github #74641 [ run ] triggered by Bot. Commit: |
|
PR_Github #74641 [ run ] completed with state
|
|
/bot run |
|
PR_Github #74655 [ run ] triggered by Bot. Commit: |
|
PR_Github #74655 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #74668 [ run ] triggered by Bot. Commit: |
|
PR_Github #74668 [ run ] completed with state
|
|
/bot run |
|
Note on the test waive in this PR CI here is blocked by a pre-existing failure on
The test(s) are waived in this PR so its CI can proceed. The same waive is being landed on |
|
PR_Github #74678 [ run ] triggered by Bot. Commit: |
|
PR_Github #74678 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
…rward calls
CapturableGuidedDecoder builds its bitmask from CUDA host functions, and
each is a callback into Python that must take the GIL. A one-engine
drafter (e.g. DFlash) runs a draft forward right after guided execution,
and that forward enters native extension calls that hold the GIL for
their whole duration (flash_attn_with_kvcache); when such a call triggers
CUDA lazy kernel-module loading, the driver performs a
context-synchronizing operation. A guided-decoder host function still
pending at that point deadlocks the rank: the device is parked at the
host-function node, the callback waits for the GIL, and the GIL holder
waits for the device. Peer ranks then spin in NCCL, wedging the whole
server with no error and no further log line. CUDA-graph warmup hits this
deterministically at the first warmup shape whose kernel instantiation is
not yet loaded (gdb evidence: pending cudaHostFuncTrampoline threads on
3/4 ranks, main threads in cudaFuncSetAttribute -> cuKernelGetFunction).
The existing set_stance("force_eager") mitigation (pytorch#163061) covers
only torch.compile kernels, not a plain extension call.
Drain at the enqueue sites: CapturableGuidedDecoder waits on
bitmask_event (recorded on the side stream after every host function) as
the last step of execute, execute_draft_batch and
rollback_rejected_batch, skipped while a CUDA graph is being captured
(host waits are illegal there, and captured host nodes do not run until
replay). Event.synchronize releases the GIL, which is exactly what lets
the pending callbacks run. Host-side ordering only; the bitmask consumer
is already ordered against the same event on the device, so results are
unchanged. No caller discipline required.
Adds CPU-only unit tests
(tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py)
asserting each enqueue site drains as its last step, that the drain is skipped under
capture or for the non-capturable decoder, and that the worker helper still executes the decoder.
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
8961297 to
c119c9e
Compare
|
PR_Github #74701 [ run ] triggered by Bot. Commit: |
|
PR_Github #74702 [ run ] triggered by Bot. Commit: |
|
PR_Github #74701 [ run ] completed with state |
|
PR_Github #74702 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #74757 [ run ] triggered by Bot. Commit: |
|
PR_Github #74757 [ run ] completed with state |
|
@brnguyen2 The dev/engineer QA review sections in the description indicate that this is a formatting change with no associated tests. Can you update the description to accurately reflect the change? |
|
Was making the drain unconditional (rather than being conditional on one-engine drafter + GIL-holding native forward) a deliberate "always-safe, some throughput cost" tradeoff, or should it be opt-in? |
|
Updated. The summary block at the top was CodeRabbit's auto-generated text from a stale pass that only saw the import reflow, so it described this as formatting-only. It now states the actual change (the drain at the three enqueue sites, skipped under graph capture) and the test coverage (the CPU-only unit test in test_guided_hostfunc_drain.py). The Description and Test Coverage sections below it were already accurate. |
|
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements. |
Deliberate. The drain is a host wait on an event the bitmask consumer already orders against on the device, so it adds no device work, and the only new cost is the CPU blocking until the host functions that were already enqueued finish. The unsafe condition is "a pending host function plus any GIL-holding call that can trigger a context sync", and the second half is not something the guided decoder can see. Gating on the drafter type would fix the case I hit and leave the same wedge for the next extension call that loads a kernel lazily. Keeping it unconditional also keeps the enqueue sites free of a mode branch: one rule (drain last, except under capture) is easier to read and to keep correct than a predicate over drafter type and attention backend that has to be revisited whenever either changes. If the wait shows up in profiles I would rather make the skip opt-in than the safety, but I have not measured a difference on the guided-decoding paths so far. |
4 participants
Dev Engineer Review
tensorrt_llm/_torch/pyexecutor/guided_decoder.py:CapturableGuidedDecodernow waits onbitmask_eventas the last step ofexecute,execute_draft_batchandrollback_rejected_batch, so the CUDA host functions that build the bitmask have finished before control returns to the caller. The wait is skipped while a CUDA graph is being captured.Event.synchronizereleases the GIL, which is what lets the pending host-function callbacks run.tensorrt_llm/_torch/speculative/interface.py: theSpecWorkerBaseguided-decoder helper routes through the draining entry points, including the skip-drafting path.QA Engineer Review
tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py(198 lines): asserts each enqueue site drains as its last step, that the drain is skipped under CUDA-graph capture and for the non-capturable decoder, and that theSpecWorkerBasehelper still executes the decoder on both the drafting and skip-drafting paths.Description
CapturableGuidedDecoderbuilds its bitmask from CUDA host functions, and each is a callback into Python that must take the GIL. A one-engine drafter (e.g. DFlash) runs a draft forward right after guided execution, and that forward enters native extension calls that hold the GIL for their whole duration (flash_attn_with_kvcache); when such a call triggers CUDA lazy kernel-module loading, the driver performs a context-synchronizing operation. A guided-decoder host function still pending at that point deadlocks the rank: the device is parked at the host-function node, the callback waits for the GIL, and the GIL holder waits for the device. Peer ranks then spin in NCCL, wedging the whole server with no error and no further log line. CUDA-graph warmup hits this deterministically at the first warmup shape whose kernel instantiation is not yet loaded (gdb evidence: pendingcudaHostFuncTrampolinethreads on 3/4 ranks, main threads incudaFuncSetAttribute->cuKernelGetFunction).The existing
set_stance("force_eager")mitigation (pytorch#163061) covers only torch.compile kernels, not a plain extension call.Fix: drain at the enqueue sites.
CapturableGuidedDecoderwaits onbitmask_event(recorded on the side stream after every host function) as the last step ofexecute,execute_draft_batchandrollback_rejected_batch, skipped while a CUDA graph is being captured (host waits are illegal there, and captured host nodes do not run until replay).Event.synchronizereleases the GIL, which is exactly what lets the pending callbacks run. Host-side ordering only; the bitmask consumer is already ordered against the same event on the device, so results are unchanged, and no caller discipline is required.Test Coverage
tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py(CPU-only): asserts each enqueue site drains as its last step, that the drain is skipped under CUDA-graph capture and for the non-capturable decoder, and that theSpecWorkerBasehelper still executes the decoder (including the skip-drafting path).PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.