Acquire the GIL in lazy init_extension

[xiuhu17]

Description

init_extension() imports Python modules via py::module_::import, which requires the GIL. Several bindings that can trigger this lazy initialization are declared with call_guard<py::gil_scoped_release> — if one of them is the first transformer_engine_torch call in the process, initialization runs without the GIL and segfaults. Typical runs don't hit this because a GIL-holding call (e.g. a quantize) usually comes first.

Deterministic repro on main:

import torch
import transformer_engine.pytorch
import transformer_engine_torch as tex

x = torch.randint(0, 255, (64, 64), dtype=torch.uint8, device="cuda")
tex.fp8_transpose(x, tex.DType.kFloat8E4M3, out=None)  # first tex call in the process

!!!!!!! Segfault encountered !!!!!!!
  ... in PyObject_Malloc
  ... in PyImport_ImportModule
  ... in transformer_engine::pytorch::init_float8_extension()
  ... in std::call_once<transformer_engine::pytorch::init_extension()::{lambda()#1}> ...
  ... in transformer_engine::pytorch::init_extension()
  ... in transformer_engine::pytorch::fp8_transpose(at::Tensor, transformer_engine::DType, std::optional<at::Tensor>)
  ... in pybind11::cpp_function::dispatcher

Running any GIL-holding tex call first (e.g. tex.quantize) and then fp8_transpose works, confirming the GIL state is the only difference.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Changes

• Acquire the GIL inside init_extension()'s std::call_once initializer before importing Python modules.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• My changes generate no new warnings

[greptile-apps]

Greptile Summary

This PR fixes a segfault that occurs when a gil_scoped_release-decorated binding (e.g. fp8_transpose) is the very first transformer_engine_torch call in the process, causing init_extension()'s lazy py::module_::import calls to run without the GIL.

• Adds pybind11::gil_scoped_acquire to the std::call_once lambda in init_extension() so that Python module imports always run with the GIL held.
• However, placing the acquire inside the lambda inverts the lock order relative to threads that already hold the GIL when they reach call_once, introducing a potential ABBA deadlock; the acquire should be moved to just before the call_once call so the order is consistently GIL → call_once internal lock.

Confidence Score: 3/5

The one-line placement of gil_scoped_acquire inside the call_once lambda trades a deterministic single-threaded segfault for a potential multi-threaded deadlock; the GIL must be acquired before call_once, not inside it.

The fix correctly identifies that py::module_::import requires the GIL, but placing gil_scoped_acquire inside the call_once lambda creates an ABBA lock ordering: a GIL-less thread can win the call_once race and then block waiting for the GIL, while a GIL-holding thread blocks waiting for the call_once internal lock. Moving the acquire to just before call_once resolves both the original segfault and the new deadlock.

Files Needing Attention: transformer_engine/pytorch/csrc/extensions/pybind.cpp — the init_extension function needs the GIL acquired outside (before) std::call_once.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Adds pybind11::gil_scoped_acquire inside the call_once lambda to fix a no-GIL segfault during lazy init, but placing the acquire inside the lambda inverts lock ordering with callers that hold the GIL, creating a potential ABBA deadlock; GIL should be acquired before call_once instead.

Sequence Diagram

sequenceDiagram
    participant A as Thread A (GIL released via call_guard)
    participant CO as std::call_once lock
    participant GIL as Python GIL
    participant B as Thread B (holds GIL, calls init_extension)

    Note over A,B: Both threads reach init_extension() for the first time concurrently
    A->>CO: enters call_once, acquires internal lock
    B->>GIL: holds GIL (normal Python thread)
    A-->>GIL: gil_scoped_acquire (BLOCKS — B holds GIL)
    B-->>CO: call_once (BLOCKS — A holds call_once lock)
    Note over A,B: DEADLOCK: A needs GIL (held by B), B needs call_once lock (held by A)
    Note over A,B: Correct fix: acquire GIL before call_once so lock order is always GIL → call_once

Loading

_{Reviews (2): Last reviewed commit: "Acquire the GIL in lazy init_extension" | Re-trigger Greptile}

[greptile-apps]

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();