diff --git a/cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp b/cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp index 0ff16e582191..41023f9f98e5 100644 --- a/cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp +++ b/cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -81,7 +81,23 @@ std::optional launchHostFunc( auto hostFuncUserData = std::make_unique(freeUserData, pyHostFunc, nb::tuple(pyArgs), nb::dict(pyKwargs)); - cudaError_t err = cudaLaunchHostFunc(stream, cudaHostFuncTrampoline, hostFuncUserData.get()); + cudaError_t err; + { + // Release the GIL only around the eager enqueue. cudaLaunchHostFunc serializes + // host-function dispatch on the stream, so in the non-capturing path it can + // busy-wait in the driver behind an in-flight callback; that callback + // (cudaHostFuncTrampoline) reacquires the GIL to run Python (e.g. xgrammar in + // guided decoding), so holding the GIL here deadlocks against it. During graph + // capture (freeUserData == false) no callback runs, so we keep the GIL to hold + // capture on a single thread. The release scope ends before the nanobind + // refcount work below, preserving the object-lifetime fix from #13251. + std::optional gilRelease; + if (freeUserData) + { + gilRelease.emplace(); + } + err = cudaLaunchHostFunc(stream, cudaHostFuncTrampoline, hostFuncUserData.get()); + } if (err != cudaSuccess) { throw std::runtime_error("Failed to launch host function."); diff --git a/tests/unittest/bindings/test_hostfunc.py b/tests/unittest/bindings/test_hostfunc.py index f208ca0ce94e..db1b1026a8f5 100644 --- a/tests/unittest/bindings/test_hostfunc.py +++ b/tests/unittest/bindings/test_hostfunc.py @@ -1,3 +1,5 @@ +import threading + import torch from tensorrt_llm._torch.hostfunc import HOSTFUNC_USER_DATA_HANDLES, hostfunc @@ -30,3 +32,52 @@ def increase(x: torch.Tensor): assert (x == 25).all().item() assert len(HOSTFUNC_USER_DATA_HANDLES) == 2 + + +def test_hostfunc_gil_release_no_deadlock(): + """Regression guard for the guided-decoding x host-func GIL deadlock. + + launch_hostfunc must not hold the GIL across cudaLaunchHostFunc. The driver + serializes host-function dispatch on a stream, so a launch issued while an + earlier callback is in flight can busy-wait in the driver. That in-flight + callback runs through cudaHostFuncTrampoline, which acquires the GIL to call + Python (e.g. xgrammar in guided decoding); if the launching thread holds the + GIL while it busy-waits, the callback can never acquire it and the two + deadlock. See cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp::launchHostFunc. + + Best-effort guard: the fixed build completes promptly; a build that holds the + GIL across the enqueue hangs and is caught by the join timeout. + """ + started = threading.Event() + gate = threading.Event() + + @hostfunc + def blocker(): + # Runs on a CUDA driver worker with the GIL held by the trampoline. + # gate.wait() drops the GIL while waiting and must reacquire it to + # return -- mimicking xgrammar reacquiring the GIL after native work. + started.set() + gate.wait(20) + + @hostfunc + def noop(): + pass + + stream = torch.cuda.Stream() + + def issue(): + with torch.cuda.stream(stream): + blocker() + noop() # second same-stream launch while blocker is in flight + torch.cuda.synchronize() + + worker = threading.Thread(target=issue, name="hostfunc-issue") + worker.start() + # Let the blocker callback begin (fixed build: the launching thread released + # the GIL so the callback can run), then release it so it can drain. + started.wait(20) + gate.set() + + worker.join(60) + assert not worker.is_alive(), ( + "launch_hostfunc deadlocked: the GIL was held across cudaLaunchHostFunc")