[symm_mem] Make the implicit MemPool path in `symm_mem.empty()` device-agnostic
- 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
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