From 71add277cb1822ad2562ac69a44339c0a6f01e13 Mon Sep 17 00:00:00 2001 From: Kirill Vasilev <290354918+vasilevklart@users.noreply.github.com> Date: Thu, 29 Oct 2026 20:00:00 +0000 Subject: [PATCH 1/2] [KV cache][Bug fix] Share a KV cache over CUDA IPC only when another worker reads it A KV cache on a LocalTransferEngine was shared with storage._share_cuda_() whatever the deployment. torch counts a shared block as held by a reader until a reader releases it, so a cache no other process ever opened went to torch's CUDA IPC limbo when its engine was dropped: its pool stayed allocated until the process exited, and at exit the limbo's destructor (CudaIPCGlobalEntities -> ~CudaIPCSentData -> CUDAGuard -> c10_cuda_check_implementation) ran after c10's CUDAKernelLaunchRegistry singleton was destroyed, reading freed memory. Under glibc's MALLOC_PERTURB_ that is a segfault at exit (the draft-verify plan-ahead GPU tests, every toy engine test); the "Producer process has been terminated before all shared CUDA tensors released" warning at exit is the same limbo. A cache whose resource spans one worker instance (needs_remote_transfer False) now gets the LocalOnlyKVTransferEngine that CPU caches already got; a shared one still goes over CUDA IPC. --- mstar/engine/resources/kv/transfer.py | 12 ++++++---- test/modular/test_kv_ipc_transfer.py | 34 +++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/mstar/engine/resources/kv/transfer.py b/mstar/engine/resources/kv/transfer.py index db41407653..5d4f6a1a41 100644 --- a/mstar/engine/resources/kv/transfer.py +++ b/mstar/engine/resources/kv/transfer.py @@ -761,9 +761,15 @@ def __init__( elif isinstance( transfer_engine_info.transfer_engine, LocalTransferEngine ): - if kv_cache.device.type == "cuda": + if not needs_remote_transfer: + # No other worker reads it, so it isn't shared: a cache shared over CUDA + # IPC and then freed waits in torch's IPC limbo for a reader to release + # it, still allocated, and the limbo's destructor makes CUDA calls at exit, + # after c10's statics are destroyed. + self._kv_transfer_engine = LocalOnlyKVTransferEngine() + elif kv_cache.device.type == "cuda": self._kv_transfer_engine = CudaIpcKVTransferEngine(kv_cache) - elif needs_remote_transfer: + else: shm_dir = transfer_engine_info.get_shm_dir() if not shm_dir: raise ValueError( @@ -775,8 +781,6 @@ def __init__( shm_dir=shm_dir, resource_key=resource_key, ) - else: - self._kv_transfer_engine = LocalOnlyKVTransferEngine() else: raise ValueError(f"Unsupported transfer engine type: {type(transfer_engine_info.transfer_engine)}") diff --git a/test/modular/test_kv_ipc_transfer.py b/test/modular/test_kv_ipc_transfer.py index 53d07c4c5a..9fc198921b 100644 --- a/test/modular/test_kv_ipc_transfer.py +++ b/test/modular/test_kv_ipc_transfer.py @@ -161,6 +161,40 @@ def test_a_read_from_another_process_lands_the_named_pages_and_tokens(): os.unlink(path) + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_a_cache_no_other_worker_reads_is_freed_with_its_manager(): + """Only a cache another worker reads is shared over CUDA IPC. A shared one that is + freed stays allocated in torch's IPC limbo (its reader holds the release) until exit, + where the limbo's destructor runs after c10's statics: every engine a process + builds and drops would keep its KV pool.""" + import gc + + from mstar.communication.tensors import LocalTransferEngine + from mstar.engine.resources.kv.transfer import ( + CudaIpcKVTransferEngine, + KVTransferManager, + LocalOnlyKVTransferEngine, + TransferEngineInfo, + ) + + info = TransferEngineInfo("h", "h", LocalTransferEngine("h")) + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + cache = KVCache(_config(max_num_pages=64, max_seq_len=512), torch.device("cuda:0")) + nbytes = cache.tensor.nbytes + manager = KVTransferManager(info, cache, needs_remote_transfer=False) + local = isinstance(manager._kv_transfer_engine, LocalOnlyKVTransferEngine) + manager.cleanup() + del manager, cache + gc.collect() + assert torch.cuda.memory_allocated() - before < nbytes + assert local + + shared = KVTransferManager(info, KVCache(_config(), torch.device("cuda:0")), needs_remote_transfer=True) + assert isinstance(shared._kv_transfer_engine, CudaIpcKVTransferEngine) + shared.cleanup() + # Pages 0 and 1 whole into 2 and 3, then tokens [0, 5) of page 2 into page 0, # every layer: the shape of a prefill -> decode handoff. Page 1 stays zero. _READER = """ From 61bc461c03f148e83e4156a02b29e5f5311a7a6d Mon Sep 17 00:00:00 2001 From: Kirill Vasilev <290354918+vasilevklart@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:14:54 -0700 Subject: [PATCH 2/2] kv transfer: shorter comment; the test checks the memory, not the class No behaviour change. The freed-with-its-manager test asserts the memory comes back rather than which engine class was picked for it. --- mstar/engine/resources/kv/transfer.py | 6 ++---- test/modular/test_kv_ipc_transfer.py | 11 +++-------- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/mstar/engine/resources/kv/transfer.py b/mstar/engine/resources/kv/transfer.py index 5d4f6a1a41..79ef2330fd 100644 --- a/mstar/engine/resources/kv/transfer.py +++ b/mstar/engine/resources/kv/transfer.py @@ -762,10 +762,8 @@ def __init__( transfer_engine_info.transfer_engine, LocalTransferEngine ): if not needs_remote_transfer: - # No other worker reads it, so it isn't shared: a cache shared over CUDA - # IPC and then freed waits in torch's IPC limbo for a reader to release - # it, still allocated, and the limbo's destructor makes CUDA calls at exit, - # after c10's statics are destroyed. + # not shared: a freed IPC-shared cache stays allocated in torch's IPC + # limbo, whose destructor makes CUDA calls after c10's statics are gone self._kv_transfer_engine = LocalOnlyKVTransferEngine() elif kv_cache.device.type == "cuda": self._kv_transfer_engine = CudaIpcKVTransferEngine(kv_cache) diff --git a/test/modular/test_kv_ipc_transfer.py b/test/modular/test_kv_ipc_transfer.py index 9fc198921b..536b42d586 100644 --- a/test/modular/test_kv_ipc_transfer.py +++ b/test/modular/test_kv_ipc_transfer.py @@ -161,20 +161,16 @@ def test_a_read_from_another_process_lands_the_named_pages_and_tokens(): os.unlink(path) - @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") def test_a_cache_no_other_worker_reads_is_freed_with_its_manager(): - """Only a cache another worker reads is shared over CUDA IPC. A shared one that is - freed stays allocated in torch's IPC limbo (its reader holds the release) until exit, - where the limbo's destructor runs after c10's statics: every engine a process - builds and drops would keep its KV pool.""" + """Only a cache another worker reads is shared over CUDA IPC; any other is freed with + its manager, not kept in torch's IPC limbo until exit.""" import gc from mstar.communication.tensors import LocalTransferEngine from mstar.engine.resources.kv.transfer import ( CudaIpcKVTransferEngine, KVTransferManager, - LocalOnlyKVTransferEngine, TransferEngineInfo, ) @@ -184,17 +180,16 @@ def test_a_cache_no_other_worker_reads_is_freed_with_its_manager(): cache = KVCache(_config(max_num_pages=64, max_seq_len=512), torch.device("cuda:0")) nbytes = cache.tensor.nbytes manager = KVTransferManager(info, cache, needs_remote_transfer=False) - local = isinstance(manager._kv_transfer_engine, LocalOnlyKVTransferEngine) manager.cleanup() del manager, cache gc.collect() assert torch.cuda.memory_allocated() - before < nbytes - assert local shared = KVTransferManager(info, KVCache(_config(), torch.device("cuda:0")), needs_remote_transfer=True) assert isinstance(shared._kv_transfer_engine, CudaIpcKVTransferEngine) shared.cleanup() + # Pages 0 and 1 whole into 2 and 3, then tokens [0, 5) of page 2 into page 0, # every layer: the shape of a prefill -> decode handoff. Page 1 stays zero. _READER = """