Segfault at interpreter exit when a torch KVCacheManager is still alive

[norx1991]

Summary

If a tpu_sync.api.torch.kv_cache_manager.KVCacheManager is still referenced when the Python interpreter shuts down, the process segfaults (exit 139). The same program exits cleanly if the manager is deleted before exit.

The manager's destructor releases its PJRT raw-buffer handles during Py_FinalizeEx. By then torch_tpu's exit teardown has already destroyed the PjRt client, so PJRT_RawBuffer_Destroy runs against a destroyed client.

Nothing in the program is wrong except that the manager lives until shutdown. That happens easily in practice:

• a manager held in a module global or a long-lived object;
• a manager pinned by any retained traceback or frame. We hit this through an unrelated library that stored an ImportError in a module global.

The work itself completes; the crash comes after. A test runner that checks the exit code reports a failure even though every test passed.

Environment

• TPU v7x, 2x2x1 slice, single host; one device used (tpu:0); Python 3.12; libtpu 0.0.47; torch 2.11.0.
• The reproducer table and native stack below: torch_tpu 0.1.1.dev20260817010134 with tpu_raiden_torch 0.0.1.dev20260906031043.
• Also seen with torch_tpu 0.1.1.dev20260913101137 and tpu-sync-torch 0.0.1.dev20260923040928. A test process segfaulted (exit -11) at interpreter exit, after its tests passed, while a manager was still alive. That run captured no native stack.
• The destructor path below is unchanged at tpu-sync 06d9a123 (2026-09-22).

Reproducer

"""Does a process that created a raiden KVCacheManager segfault at exit?"""
import os, sys, time
import torch
import torch_tpu  # noqa: F401
from tpu_sync.api.torch.kv_cache_manager import KVCacheManager

mode = sys.argv[1] if len(sys.argv) > 1 else "transfer"
os.environ["RAIDEN_ENABLE_ASYNC_DISPATCH"] = "1"
buf = torch.ones((8, 128, 2, 2, 128), dtype=torch.bfloat16).to("tpu:0")
if mode != "tensor_only":
    mgr = KVCacheManager([[buf]], local_control_port=0, host_blocks_to_allocate=8,
                         unsafe_skip_buffer_lock=True)
    if mode in ("transfer", "transfer_del"):
        f = mgr.d2h([1], [0], [1])
        while not f.is_ready(): time.sleep(0.001)
        f = mgr.h2d([0], [4], [1])
        while not f.is_ready(): time.sleep(0.001)
        torch.tpu.synchronize()
    if mode in ("create_del", "transfer_del"):
        del mgr
print(f"mode={mode} done, exiting", flush=True)

mode	what it does	exit code
tensor_only	TPU tensor only, no manager	0
create	manager created, alive at exit	139
create_del	manager created, then del	0
transfer	manager plus one D2H and one H2D, alive at exit	139
transfer_del	same, then del	0

The DMAs don't matter: holding the manager until exit is enough. Managers destroyed mid-process are fine.

Stack

*** SIGSEGV (@(nil)), si_code=1 ***
    @  FailureSignalHandler()
    @  tsl::AsyncValue::MakeTypeInfo<>()::{lambda()#1}::__invoke()
    @  tsl::AsyncValue::Destroy()
    @  xla::TpuRawBuffer::~TpuRawBuffer()
    @  pjrt::PJRT_RawBuffer_Destroy()
    @  pjrt::PjRtCApiRawBuffer_Destroy()
    @  std::_Sp_counted_ptr_inplace<>::_M_dispose()
    @  raiden::RaidenBufferHandle::~RaidenBufferHandle()
    @  tpu_raiden::kv_cache::KVCacheManagerBase::~KVCacheManagerBase()
    @  tpu_raiden::KVCacheManagerWithTransfer::~KVCacheManagerWithTransfer()
    @  tpu_raiden::torch::TorchKVCacheManager::~TorchKVCacheManager()
    @  tpu_raiden::torch::KVCacheManager::~KVCacheManager()
    @  (python)
    @  _PyModule_ClearDict
    @  Py_FinalizeEx
    @  Py_RunMain

torch_tpu's exit-time output (compilation-cache stats such as num_cache_hits=..., peak_compilation_memory_bytes=...) is printed immediately before the crash. So the client teardown has already run when the manager is destroyed.

Analysis

• KVCacheManagerBase::~KVCacheManagerBase() (tpu_sync/kv_cache/kv_cache_manager_base.cc:513) ends with layers_.clear(). That destroys each layer's RaidenBufferHandle (tpu_sync/core/raw_transfer_core.h:133).
• Each handle holds std::shared_ptr<RawBufferHolder> c_hold. ~RawBufferHolder() (raw_transfer_core.h:55) calls pjrt::PjRtCApiRawBuffer_Destroy unconditionally.
• At interpreter shutdown, Python clears module dicts (_PyModule_ClearDict) after the atexit handlers have run. torch_tpu tears down its PjRt client in its exit path, so the last reference to the manager drops after that. The raw-buffer destroy then touches the client's freed state.
• Nothing orders the manager's destruction before the client's teardown, and the manager holds nothing that keeps the client alive.

Possible fixes

1. Release live managers before the client goes away (Python side).
   • Track managers in a weakref.WeakSet, and have an atexit handler release their native state (for example, a new close() that drops _impl).
   • Register the handler on the first KVCacheManager construction, not at import. The PjRt client must already exist by then, because the manager needs device buffers. atexit runs handlers in reverse order of registration, so this one would run before torch_tpu's teardown handler, assuming torch_tpu registers its handler when the client is created.
   • An explicit, documented close() would also help callers that want deterministic teardown.
2. Don't touch PJRT after the client is gone (C++ side).
   • When the destructor runs during finalization (Py_IsFinalizing() in the torch binding), or after a client-destroyed flag has been set, release the RawBufferHolders without calling PJRT_RawBuffer_Destroy.
   • The process is exiting anyway, so leaking these handles is harmless.
3. Tie the manager's lifetime to the client: hold a reference that keeps the PjRt client alive until the manager is gone. This only works if torch_tpu's teardown respects outstanding references, so it probably also needs a change in torch_tpu.

Option 1, with option 2 as a backstop, seems the least invasive.

Workaround for users

Delete every KVCacheManager, and anything that owns one, before the interpreter exits, and make sure no retained traceback or frame still references it. Calling os._exit after the work is done also avoids the crash, but it skips all other cleanup.