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 binding registered with call_guard<gil_scoped_release> is the first transformer_engine_torch call in the process: init_extension() would reach py::module_::import without holding the GIL, crashing inside PyImport_ImportModule. The fix acquires the GIL before std::call_once, ensuring the same lock order (GIL → call_once internal lock) on every entry path and eliminating both the segfault and the lock-inversion deadlock that would occur if the acquire were placed inside the lambda.

• pybind11::gil_scoped_acquire gil is placed before std::call_once in init_extension(), so all threads — whether or not they held the GIL on entry — always acquire it before competing for the call_once lock.
• After initialization, callers that ran under call_guard<gil_scoped_release> will briefly re-acquire and release the GIL on every subsequent init_extension() call; an atomic fast-path guard (see inline suggestion) can eliminate that overhead.

Confidence Score: 5/5

Safe to merge — the one-line change correctly acquires the GIL before call_once, fixing the reported segfault and maintaining a consistent lock order across all callers.

The change is minimal and mechanically correct: placing gil_scoped_acquire before call_once ensures every thread establishes the same lock order (GIL first, then the call_once internal mutex), eliminating the original segfault. No data races or ordering inversions are introduced.

Files Needing Attention: No files require special attention; the single changed function in pybind.cpp is straightforward.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Adds pybind11::gil_scoped_acquire before std::call_once in init_extension() to fix a segfault when a GIL-released binding is the first tex call; the GIL-before-call_once ordering eliminates the lock inversion risk raised in prior review.

Sequence Diagram

sequenceDiagram
    participant PY as Python caller
    participant DIS as pybind11 dispatcher
    participant FN as fp8_transpose
    participant INIT as init_extension
    participant ONCE as call_once initializer

    PY->>DIS: "tex.fp8_transpose(...)"
    note over DIS: "call_guard drops GIL"
    DIS->>FN: "(GIL released)"
    FN->>INIT: "init_extension()"
    INIT->>INIT: "gil_scoped_acquire: re-acquire GIL"
    INIT->>ONCE: "std::call_once(flag, lambda)"
    alt "First call: flag not set"
        ONCE->>ONCE: "import Python modules (needs GIL)"
        ONCE-->>INIT: "flag set"
    else "Subsequent calls: flag already set"
        ONCE-->>INIT: "no-op"
    end
    INIT->>INIT: "~gil_scoped_acquire: release GIL"
    INIT-->>FN: return
    FN-->>PY: result

Loading

_{Reviews (3): Last reviewed commit: "Update transformer_engine/pytorch/csrc/e..." | Re-trigger Greptile}

[ptrendx]

Actually, I agree with the Greptile review. @xiuhu17, could you make the suggested change?

Waiting for the suggested changes to be implemented.

[ptrendx]

/te-ci pytorch

[xiuhu17]

Actually, I agree with the Greptile review. @xiuhu17, could you make the suggested change?

Hi @ptrendx, It seems that the ci fail is not related to our change?

[ptrendx]

Yes, merged. Thank you for your contribution:-)!

[xiuhu17]

@ptrendx Thanks for reviewing!