Skip to content

Acquire the GIL in lazy init_extension #3255

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xiuhu17 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Acquire the GIL in lazy init_extension #3255

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions transformer_engine/pytorch/csrc/extensions/pybind.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,7 @@ void init_grouped_tensor_extension() {

void init_extension() {
std::call_once(extension_init_flag, []() {
pybind11::gil_scoped_acquire gil;
init_float8_extension();
Comment on lines 122 to 125

@greptile-apps greptile-apps Bot Jul 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Placing gil_scoped_acquire inside the call_once lambda inverts the lock order relative to callers that still hold the GIL, creating an ABBA deadlock. Consider two threads both hitting an uninitialised tex entry point simultaneously: Thread A (GIL released via call_guard<gil_scoped_release>) wins the call_once race and blocks on gil_scoped_acquire; Thread B (holds the GIL, e.g. calling quantize) reaches call_once and blocks waiting for Thread A's lambda to complete. Thread A owns the call_once internal lock and needs the GIL; Thread B owns the GIL and needs the call_once lock — neither can proceed.

The fix is to acquire the GIL before call_once so that every caller establishes the same ordering (GIL → call_once lock), eliminating the inversion. After the one-time init completes, subsequent calls to init_extension() will acquire the GIL briefly and return immediately from call_once; pybind11::gil_scoped_acquire is a no-op when the calling thread already holds the GIL, so callers that never released it are unaffected.

Suggested change
void init_extension() {
std::call_once(extension_init_flag, []() {
pybind11::gil_scoped_acquire gil;
init_float8_extension();
void init_extension() {
pybind11::gil_scoped_acquire gil;
std::call_once(extension_init_flag, []() {
init_float8_extension();

init_mxfp8_extension();
init_float8blockwise_extension();
Expand Down