[GH-ISSUE #799] [HybridEP] Fabric (MNNVL) allocator leaks device memory on every buffer reallocation #484

New issue

Open

opened 2026-10-27 21:02:45 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:02:45 +00:00

Owner

Originally created by @jerryhong21 on GitHub (Oct 26, 2026).
Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/799

Summary

On systems where HybridEP uses the fabric (MNNVL) allocator (for example GB200/GB300 NVL72), physical device memory is never returned when a HybridEP buffer is freed. Each time HybridEPBuffer::update_buffer grows the buffer, the entire previous buffer set leaks. With variable-length batches, the leaks keep adding up until cuMemCreate fails.

Symptom

We hit this in MoE training with variable-length, long-context batches on GB300 NVL72, using HybridEP through Megatron-LM's flex dispatcher:

Device memory on one rank climbed stepwise from ~242 GiB to ~281 GiB over ~2h and never came back down.
PyTorch's allocator stats did not account for the growth, because this memory is allocated outside PyTorch.
The run then aborted mid-step:

CU error encountered at: file=.../csrc/hybrid_ep/allocator/allocator.cu line=91,
call='cuMemCreate(&handle, size, &fabric_prop_, 0)' Reason=out of memory.

Root cause

In csrc/hybrid_ep/allocator/allocator.cu, the fabric path never releases three handle references:

allocate never releases the handle from cuMemCreate (:91).
get_handle never releases the reference taken by cuMemRetainAllocationHandle before export (:124).
open_handle never releases the handle from cuMemImportFromShareableHandle (:139).

free retains and releases once, so its net effect on the count is zero. close_handle only unmaps. The CUDA driver frees physical memory only after every mapping is unmapped and every handle reference is released, so this memory stays allocated. Peers' imports also keep each exporter's memory pinned.

update_buffer calls release_buffer() and then allocate_buffer() whenever a call needs more capacity than the current reservation (hybrid_ep.cu, grow_buffer_config). With variable-length inputs, every new tokens-per-rank high (in alignment steps) therefore leaks a whole buffer set. Workloads where every step has the same shape allocate once and are not affected.

The cudaMalloc/IPC path, used when fabric is unavailable, is not affected. The legacy SharedMemoryAllocator in csrc/deep_ep.cpp has the same three sites.

Workarounds

Size the buffer up front for the maximum max_num_of_tokens_per_rank, so it never reallocates.
Or set USE_MNNVL=0 when the EP group fits within one NVLink node, so the cudaMalloc/IPC path is used.

Fix

PR #798 adds the missing cuMemRelease calls after map, export and import.

Originally created by @jerryhong21 on GitHub (Oct 26, 2026). Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/799 ## Summary On systems where HybridEP uses the fabric (MNNVL) allocator (for example GB200/GB300 NVL72), physical device memory is never returned when a HybridEP buffer is freed. Each time `HybridEPBuffer::update_buffer` grows the buffer, the entire previous buffer set leaks. With variable-length batches, the leaks keep adding up until `cuMemCreate` fails. ## Symptom We hit this in MoE training with variable-length, long-context batches on GB300 NVL72, using HybridEP through Megatron-LM's flex dispatcher: - Device memory on one rank climbed stepwise from ~242 GiB to ~281 GiB over ~2h and never came back down. - PyTorch's allocator stats did not account for the growth, because this memory is allocated outside PyTorch. - The run then aborted mid-step: ``` CU error encountered at: file=.../csrc/hybrid_ep/allocator/allocator.cu line=91, call='cuMemCreate(&handle, size, &fabric_prop_, 0)' Reason=out of memory. ``` ## Root cause In `csrc/hybrid_ep/allocator/allocator.cu`, the fabric path never releases three handle references: - `allocate` never releases the handle from `cuMemCreate` (`:91`). - `get_handle` never releases the reference taken by `cuMemRetainAllocationHandle` before export (`:124`). - `open_handle` never releases the handle from `cuMemImportFromShareableHandle` (`:139`). `free` retains and releases once, so its net effect on the count is zero. `close_handle` only unmaps. The CUDA driver frees physical memory only after every mapping is unmapped *and* every handle reference is released, so this memory stays allocated. Peers' imports also keep each exporter's memory pinned. `update_buffer` calls `release_buffer()` and then `allocate_buffer()` whenever a call needs more capacity than the current reservation (`hybrid_ep.cu`, `grow_buffer_config`). With variable-length inputs, every new tokens-per-rank high (in alignment steps) therefore leaks a whole buffer set. Workloads where every step has the same shape allocate once and are not affected. The `cudaMalloc`/IPC path, used when fabric is unavailable, is not affected. The legacy `SharedMemoryAllocator` in `csrc/deep_ep.cpp` has the same three sites. ## Workarounds - Size the buffer up front for the maximum `max_num_of_tokens_per_rank`, so it never reallocates. - Or set `USE_MNNVL=0` when the EP group fits within one NVLink node, so the `cudaMalloc`/IPC path is used. ## Fix PR #798 adds the missing `cuMemRelease` calls after map, export and import.

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

deepseek-ai/DeepEP#484

No description provided.

Rows
Columns