[PR #798] [CLOSED] [HybridEP] Release fabric memory handles after map/export/import #792

Closed
opened 2026-10-27 21:06:25 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/798
Author: @jerryhong21
Created: 10/26/2026
Status: ❌ Closed

Base: hybrid-ep ← Head: hybridep-fabric-handle-release


📝 Commits (1)

  • 14360b3 [HybridEP] Release fabric memory handles after map/export/import

📊 Changes

2 files changed (+8 additions, -0 deletions)

View changed files

📝 csrc/deep_ep.cpp (+4 -0)
📝 csrc/hybrid_ep/allocator/allocator.cu (+4 -0)

📄 Description

Problem

On the fabric (MNNVL) path, ExtendedMemoryAllocator never releases three handle references:

  • allocate: the handle returned by cuMemCreate (allocator.cu:91).
  • get_handle: the reference taken by cuMemRetainAllocationHandle before export (:124).
  • open_handle: the handle returned by cuMemImportFromShareableHandle (:139).

free retains a reference and releases it 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 is never returned.

HybridEPBuffer::update_buffer calls release_buffer() and then allocate_buffer() whenever a call needs more capacity than the current reservation. With variable-length batches, every new tokens-per-rank high therefore leaks the entire previous buffer set on the exporting rank, and every peer's import keeps it pinned. On GB300 NVL72 with variable-length long-context MoE training, we saw device memory climb stepwise over a ~2h run, never released, until cuMemCreate failed with out of memory at allocator.cu:91 mid-step. PyTorch's allocator stats did not show the growth, because this memory is allocated outside PyTorch.

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.

Fix

Release each reference once the mapping or export holds what it needs:

  • after cuMemMap in allocate;
  • after cuMemExportToShareableHandle in get_handle;
  • after cuMemMap in open_handle.

The mapping keeps the memory alive, so free and close_handle now actually return it. csrc/deep_ep.cpp gets the same fix.

Testing

A GB300 regrow-loop measurement (device memory per reallocation, before vs. after) is in progress. I'll post the numbers here when it finishes.


🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/798 **Author:** [@jerryhong21](https://github.com/jerryhong21) **Created:** 10/26/2026 **Status:** ❌ Closed **Base:** `hybrid-ep` ← **Head:** `hybridep-fabric-handle-release` --- ### 📝 Commits (1) - [`14360b3`](https://github.com/deepseek-ai/DeepEP/commit/14360b3593dcc0a38b723c02b4479d97933b6d36) [HybridEP] Release fabric memory handles after map/export/import ### 📊 Changes **2 files changed** (+8 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `csrc/deep_ep.cpp` (+4 -0) 📝 `csrc/hybrid_ep/allocator/allocator.cu` (+4 -0) </details> ### 📄 Description ## Problem On the fabric (MNNVL) path, `ExtendedMemoryAllocator` never releases three handle references: - `allocate`: the handle returned by `cuMemCreate` (`allocator.cu:91`). - `get_handle`: the reference taken by `cuMemRetainAllocationHandle` before export (`:124`). - `open_handle`: the handle returned by `cuMemImportFromShareableHandle` (`:139`). `free` retains a reference and releases it 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 is never returned. `HybridEPBuffer::update_buffer` calls `release_buffer()` and then `allocate_buffer()` whenever a call needs more capacity than the current reservation. With variable-length batches, every new tokens-per-rank high therefore leaks the entire previous buffer set on the exporting rank, and every peer's import keeps it pinned. On GB300 NVL72 with variable-length long-context MoE training, we saw device memory climb stepwise over a ~2h run, never released, until `cuMemCreate` failed with `out of memory` at `allocator.cu:91` mid-step. PyTorch's allocator stats did not show the growth, because this memory is allocated outside PyTorch. 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. ## Fix Release each reference once the mapping or export holds what it needs: - after `cuMemMap` in `allocate`; - after `cuMemExportToShareableHandle` in `get_handle`; - after `cuMemMap` in `open_handle`. The mapping keeps the memory alive, so `free` and `close_handle` now actually return it. `csrc/deep_ep.cpp` gets the same fix. ## Testing A GB300 regrow-loop measurement (device memory per reallocation, before vs. after) is in progress. I'll post the numbers here when it finishes. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
bsakaguchi 2026-10-27 21:06:25 +00:00
Sign in to join this conversation.
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#792
No description provided.