[sglang] CUDA IPC 핸들 하나를 여러 피어가 열면 생기는 일: SGLang diffusion all-to-all 메모리 누수 수정기

PR 링크: sgl-project/sglang#50321 상태: Merged | 변경: +27 / -3

들어가며

SGLang의 multimodal diffusion 런타임은 Ulysses sequence parallel에서 all-to-all 통신을 NCCL 대신 CUDA IPC 기반 직접 메모리 매핑으로 처리하는 경로(ipc_a2a_multi.py)를 갖고 있습니다. Ulysses degree가 2보다 크고 head-group pipeline을 쓰는 경우, 각 rank는 자신의 버퍼를 IPC 핸들로 export하고 다른 rank들은 그 핸들을 열어 상대 GPU 메모리에 직접 씁니다.

문제는 요청 geometry가 바뀔 때마다 GPU 메모리가 새는 것이었습니다. 8x B300, Ulysses 8 환경에서 요청마다 sequence length가 다른 video DiT를 돌리면 요청당 약 2GB씩 live memory가 증가했고, 약 50번째 요청에서 OOM이 발생했습니다. 고정 shape 워크로드도 느리게 누수됐습니다.

원인은 PyTorch CUDA IPC의 reference counting 방식에 있었습니다.

• reduce_tensor(t)로 만든 하나의 핸들을 world - 1개 피어가 모두 열었습니다.
• torch는 producer 쪽에서 핸들 하나당 하나의 IPC reference count를 유지하고, 그 count가 0으로 돌아와야 producer의 블록을 해제합니다.
• 여러 consumer가 같은 하나의 count를 각자 release해도 count 관리가 consumer 수와 맞지 않아 0으로 돌아오지 않았고, 결국 버려진 버퍼는 영원히 해제되지 않았습니다.

누수 지점은 두 곳이었습니다. 새 plan이 만들어질 때 버려지는 _PipelinePool, 그리고 LRU에서 evict되는 staging buffer입니다.

코드 분석

모든 변경은 ipc_a2a_multi.py 한 파일에 있지만, 역할별로 세 부분으로 나눠 보겠습니다.

1. _share: 피어마다 별도 핸들 발급

Before

handles: list = [None] * self.world
dist.all_gather_object(handles, reduce_tensor(t), group=group)
dev = torch.cuda.current_device()
out = []
for member, (fn, args) in enumerate(handles):
    if member == self.rank:
        out.append(t)
        continue

각 rank는 reduce_tensor(t)를 한 번만 호출하고, 그 결과를 all-gather 했습니다. 즉 rank 0의 핸들 하나를 rank 1~7이 모두 열었습니다.

After

mine = [
    None if member == self.rank else reduce_tensor(t)
    for member in range(self.world)
]
handles: list = [None] * self.world
dist.all_gather_object(handles, mine, group=group)
dev = torch.cuda.current_device()
out = []
for member, theirs in enumerate(handles):
    if member == self.rank:
        out.append(t)
        continue
    fn, args = theirs[self.rank]

이제 각 rank는 피어 수만큼 reduce_tensor를 호출해 world 길이의 리스트를 만듭니다(자기 자신 자리는 None). all-gather 후 handles[member]는 member가 만든 리스트이고, 수신 측은 theirs[self.rank] — 즉 자신만을 위해 발급된 핸들만 엽니다. 핸들과 consumer가 1:1로 대응되므로 각 consumer가 닫을 때 해당 핸들의 count가 정확히 0으로 돌아옵니다.

docstring에도 이 의도가 명시되었습니다.

"""...
Each peer gets its own handle: torch counts one reference per handle and
frees `t` only once that count is back to zero, so one handle opened by
world - 1 peers never frees it (every dropped pool or staging buffer
stayed allocated)."""

비용은 all-gather 페이로드가 world - 1배 커지는 것뿐인데, 이는 pool/staging 생성 시 한 번만 발생하는 control-plane 비용이라 무시할 만합니다.

2. _release_shared: 그룹 단위로 해제를 확정

핸들을 고쳐도 해제는 즉시 일어나지 않습니다. producer의 블록은 모든 피어가 매핑을 닫은 뒤, 그리고 torch.cuda.ipc_collect()가 IPC limbo를 정리할 때 비로소 해제됩니다. 이를 위한 헬퍼가 추가되었습니다.

def _release_shared(self) -> None:
    """Free the buffers just dropped on every rank; a group collective.

    A shared block outlives its last local reference until each peer has
    closed its mapping, and then only when the IPC limbo is collected.
    """
    torch.cuda.synchronize()
    dist.barrier(group=self.group)
    torch.cuda.ipc_collect()

순서가 핵심입니다.

1. synchronize(): 이 rank에서 해당 매핑을 대상으로 한 queued kernel/copy가 모두 끝나게 합니다.
2. barrier: 모든 rank가 자신의 참조를 놓은 시점까지 기다립니다. 그래야 producer 입장에서 count가 0이 된 상태가 보장됩니다.
3. ipc_collect(): limbo에 남은 블록을 실제로 해제합니다.

barrier가 없다면 일부 피어가 아직 매핑을 쥔 상태에서 collect가 호출되어 아무 것도 해제되지 않을 수 있습니다. PR 설명에 "torch.cuda.ipc_collect() alone freed nothing"이라고 적힌 것도 같은 맥락입니다 — 핸들 구조가 잘못되어 있으면 collect만으로는 소용이 없습니다.

3. 호출 지점: evict와 pool 재생성

staging LRU eviction

            # queued copies may still target the evicted peer mappings
            torch.cuda.synchronize()
            self.staging.popitem(last=False)
+           self._release_shared()
        self.staging[key] = value

pipeline pool 재구성

                torch.cuda.synchronize()
                del self.pools[q.dtype]
                pool = None
+               self._release_shared()
            if check is not None:
                reference = sequential()
            pool = self.pools[q.dtype] = _PipelinePool(self, q.dtype, *sizes)

_release_shared는 barrier를 포함하는 collective이므로 모든 rank가 같은 순서로 호출해야 deadlock이 없습니다. 두 지점 모두 모든 rank에서 동일한 순서로 도달하는 경로이기 때문에 안전합니다. 또한 새 _PipelinePool을 할당하기 직전에 해제하므로, 피크 메모리 시점에 old/new 버퍼가 공존하지 않습니다.

마지막으로 drop_staging()에도 정리 호출이 추가되었습니다.

        self.staging.clear()
        self.pools.clear()
+       if torch.cuda.is_available():
+           torch.cuda.ipc_collect()

참고로 2-rank 경로(ipc_a2a.py)는 핸들당 consumer가 원래 하나뿐이라 영향이 없습니다.

왜 이게 좋은가

수치

rank 0에서 torch.cuda.memory._snapshot 기반 allocation-stack census를 돌린 결과, _PipelinePool.__init__을 통해 할당된 live bytes가 14개 요청 동안 1.31GB → 32.3GB로 증가했고 이는 전체 증가분의 약 95%였습니다. 원인을 데이터로 정확히 지목한 셈입니다.

workload	before	after
가변 geometry video 요청 40개	108.9 → 194.3 GB (~50 요청에서 OOM)	108.9 → 112.1 GB, 3번째 요청부터 flat
고정 shape video 요청 20개	107.6 → 111.8 GB	107.6 → 109.3 GB

• standalone 8-process repro: 공유 핸들 하나일 때 iteration당 1GiB 누수, 피어별 핸들일 때 누수 0.
• latency: 평균 17.19s → 17.70s로 큰 차이 없음(fix 측 run은 더 무거운 allocation recording hook을 달고 있었음).
• 출력은 mp4 md5 기준 40/40, 20/20 byte-identical.

일반적 교훈

1. IPC 핸들은 "공유 포인터"가 아니라 "refcount 계약"이다. reduce_tensor 같은 serialization 결과를 브로드캐스트해 여러 프로세스가 재사용하면 refcount 시맨틱이 깨질 수 있습니다. consumer 하나당 핸들 하나를 원칙으로 삼는 것이 안전합니다.
2. 크로스 프로세스 해제는 collective다. 로컬에서 del했다고 메모리가 반환되지 않습니다. sync → barrier → collect 순서로 "모두가 놓았다"를 확정해야 합니다. 그리고 collective는 모든 rank가 결정론적으로 같은 순서로 도달하는 지점에서만 호출해야 합니다.
3. 누수 조사는 allocation stack으로. torch.cuda.memory._snapshot으로 할당 지점별 live bytes를 집계하면 "어디서 새는지"가 바로 드러납니다. 추측성 ipc_collect() 추가만으로는 해결되지 않았다는 점이 이를 잘 보여줍니다.
4. 정답성 검증을 함께 하라. 메모리 수정이 출력을 바꾸지 않았음을 md5로 확인한 것은 통신 레이어 변경에서 특히 좋은 습관입니다.

리뷰 과정에서는 코드 자체에 대한 이견 없이, CI의 PR 제목 형식 검사([Type] scope: title) 문제만 지적되어 제목이 현재 형태로 수정되었습니다.

참고 자료

• torch.cuda.ipc_collect — CUDA IPC로 해제된 메모리를 강제로 회수하는 함수
• torch.cuda.synchronize — 해제 전 queued 커널 완료 보장에 사용
• torch.distributed — barrier, all_gather_object 등 collective API
• Multiprocessing best practices - CUDA in multiprocessing — CUDA 텐서 공유와 refcount 관련 공식 노트
• Understanding CUDA Memory Usage — torch.cuda.memory._snapshot을 이용한 메모리 분석

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.