pytorch / pytorch/pytorch

[symm_mem] Make the implicit MemPool path in `symm_mem.empty()` device-agnostic

Open
#192,488 1 comment 1 reaction 0 assignees View on GitHub
bot-triaged oncall: distributed oncall: distributed infra
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🚀 The feature, motivation and pitch

`symm_mem.empty()` allocates through an implicit MemPool by default (#172292), but the path is gated to `device.type == "cuda"` because it calls `torch.cuda.use_mem_pool` directly. There is a TODO noting the path can be made device-agnostic once `use_mem_pool` is elevated to `torch.accelerator`.

The C++ side is already device-generic: `MemPoolAllocatorMap` in `SymmetricMemory.cpp` keys mempool allocators by device type, and `register_mempool_allocator()` is `TORCH_API`, so an out-of-tree PrivateUse1 backend (e.g. torch_npu) can already register its allocator but it cannot reach the pool path from `empty()` due to the CUDA-only gate. This came up concretely in #190629, where the guidance was to use the MemPool path instead of a new allocator hook, but that path is not reachable for non-CUDA backends today.

Proposal (interim step that does not wait for the full `torch.accelerator.MemPool` generalization in #167210):

1. replace the `device.type == "cuda"` gate with a check for whether the device type has a registered mempool allocator.
2. dispatch via `torch.get_device_module(device).use_mem_pool`, so XPU and out-of-tree backends that ship a device module can take the pool path.
3. when the accelerator-level API from #167210 lands, this collapses to `torch.accelerator.use_mem_pool` with no user-facing change.

I'm happy to put up a PR if this direction sounds right.

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @aditvenk @weifengpy @kapilsh @kwen2501 @ngimel

### Alternatives

Wait for RFC #167210 to land `torch.accelerator.MemPool` / `use_mem_pool` and only then generalize the gate. That is cleaner long-term but leaves PrivateUse1 backends without a supported `symm_mem.empty()` path in the interim.

### Alternatives

_No response_

### Additional context

_No response_

Contributor guide

Open the contributing guide

Research direction

Start by tracing symm_mem.empty() and its device.type == "cuda" gate, then inspect MemPoolAllocatorMap in SymmetricMemory.cpp and the registered allocator check. Verify how torch.get_device_module(device).use_mem_pool is exposed for device modules. Done means registered non-CUDA and PrivateUse1 devices can reach the implicit MemPool path without changing CUDA behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
backend
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.