[vllm] vLLM, Diffusion-Gemma 샘플러 메모리 최적화: 요청 기반 타일링으로 OOM 문제 해결

PR 링크: vllm-project/vllm#45672 상태: Merged | 변경: +104 / -90

들어가며

최근 대규모 언어 모델(LLM)의 발전과 함께 이미지 생성 모델과의 결합이 활발히 이루어지고 있습니다. Diffusion-Gemma와 같은 모델은 텍스트 설명을 기반으로 고품질 이미지를 생성하는 데 탁월한 성능을 보입니다. 하지만 이러한 모델을 효율적으로 서빙하는 것은 복잡한 기술적 과제를 동반합니다. 특히, vLLM은 LLM 추론 속도 향상을 위한 혁신적인 기술들을 선보여왔지만, 특정 모델 아키텍처에서는 예상치 못한 메모리 병목 현상이 발생하곤 합니다.

이번 글에서는 vLLM 프로젝트의 GitHub Pull Request(PR) 중 하나를 심층 분석하여, Diffusion-Gemma 모델의 샘플링 과정에서 발생하는 치명적인 Out-of-Memory(OOM) 문제를 어떻게 해결했는지 살펴보겠습니다. 이 PR은 max_num_seqs (동시 요청 수)가 증가함에 따라 기하급수적으로 늘어나는 임시(transient) 로그 확률(logit) 버퍼 메모리를 효과적으로 관리하는 데 초점을 맞추고 있습니다. 기존 방식으로는 모델 가중치나 KV 캐시보다 먼저 OOM이 발생하는 심각한 문제가 있었으나, 본 PR은 '요청 기반 타일링(request-tiled logits)' 기법을 도입하여 이 문제를 해결했습니다.

코드 변경 사항을 상세히 살펴보고, 이 최적화가 왜 효과적인지, 그리고 어떤 일반적인 교훈을 얻을 수 있는지 알아보겠습니다.

코드 분석: 파일별 변경 사항

vllm/model_executor/models/diffusion_gemma.py

이 PR의 핵심 로직은 diffusion_gemma.py 파일 내의 샘플링 관련 함수 수정에 집중되어 있습니다. 특히 _compiled_sample_step 함수와 이를 호출하는 __call__ 메소드에서 중요한 변경이 이루어졌습니다.

1. _compiled_sample_step 함수의 역할 변경 및 sampled/num_sampled 초기화 로직 수정

기존에는 _compiled_sample_step 함수 내부에서 sampled와 num_sampled 버퍼를 매번 초기화했습니다. 이는 각 호출마다 독립적인 계산을 보장하려는 의도였을 수 있으나, 타일링(tiling) 방식과 결합될 때 문제가 발생할 수 있습니다.

Before:

-    # Clear outputs so prefill / non-decode slots report 0 (decode slots are
-    # overwritten below).
-    sampled.zero_()
-    num_sampled.zero_()

After: (해당 로직이 __call__ 메소드로 이동)

_compiled_sample_step 함수 자체에서는 더 이상 sampled와 num_sampled를 초기화하지 않습니다. 이 변경은 __call__ 메소드에서 타일링된 반복문이 시작되기 전에 단 한 번만 초기화하도록 로직을 옮김으로써, 각 타일이 이전 타일의 계산 결과를 덮어쓰는 것을 방지하기 위함입니다.

2. __call__ 메소드: 요청 기반 타일링 도입 및 메모리 관리

가장 핵심적인 변경은 __call__ 메소드 내에서 샘플링 로직을 처리하는 방식입니다. 기존에는 단일 _compiled_sample_step 호출로 모든 디코딩 슬롯(decode slots)을 처리했지만, 이제는 GPU의 가용 메모리를 기반으로 동적으로 결정된 그룹 크기(group)로 요청들을 나누어 여러 번의 _compiled_sample_step 호출을 수행합니다.

Before (단일 호출):

-        # --- Single compiled call: temp → sample → probs → post-process ---
-        scaled = _compiled_sample_step(
-            logits,
-            decode_slots,
-            decode_idx,
-            all_slots,
-            valid_canvas_len,
-            # State
-            states.canvas,
-            states.argmax_canvas,
-            states.step,
-            states.is_encoder_phase,
-            states.confident,
-            states.self_conditioning_embeds,
-            self.embed_weight,
-            self.normalizer,
-            states.accepted_canvas_history,
-            states.accepted_canvas_history_len,
-            # Output
-            sampled,
-            num_sampled,
-            self.req_states.draft_tokens,
-            # Config
-            max_denoising_steps=float(states.max_denoising_steps),
-            t_min=self.t_min,
-            t_max=self.t_max,
-            confidence_threshold=self.confidence_threshold,
-            vocab_size=self.vocab_size,
-            CL=self.canvas_length,
-            ST=states.stability_threshold,
-            entropy_bound=self.entropy_bound,
-            sc_vocab_start=self.sc_vocab_start,
-            sc_vocab_end=self.sc_vocab_end,
-            tp_size=self.tp_size,
-            tp_group_name=self.tp_group_name,
-        )

After (타일링 루프 도입):

+        # Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in
+        # _compiled_sample_step keeps several live [group * CL, vocab] copies, so
+        # size each tile to a fraction of free memory to bound the transient at
+        # high concurrency. Tiling is bit-identical to a single pass.
+        group = max(num_decode, 1)
+        if num_decode > 0:
+            free, _ = current_platform.mem_get_info()
+            # ~10 transient fp32 copies of [group * CL, vocab] inside the step
+            # (eager peaks at ~8; pad for allocator overhead and small tensors).
+            bytes_per_req = CL * self.vocab_size * 4 * 10
+            budget = int(free * 0.5) // max(bytes_per_req, 1)
+            group = max(1, min(num_decode, budget))
+
+        for start_req in range(0, num_decode, group):
+            end_req = min(start_req + group, num_decode)
+            tile = slice(start_req, end_req)
+            tile_slots = decode_slots[tile]
+
+            scaled = _compiled_sample_step(
+                logits[start_req * CL : end_req * CL],
+                tile_slots,
+                decode_idx[tile],
+                all_slots,
+                valid_canvas_len[tile],
+                # State
+                states.canvas,
+                states.argmax_canvas,
+                states.step,
+                states.is_encoder_phase,
+                states.confident,
+                states.self_conditioning_embeds,
+                self.embed_weight,
+                self.normalizer,
+                states.accepted_canvas_history,
+                states.accepted_canvas_history_len,
+                # Output
+                sampled,
+                num_sampled,
+                self.req_states.draft_tokens,
+                # Config
+                max_denoising_steps=float(states.max_denoising_steps),
+                t_min=self.t_min,
+                t_max=self.t_max,
+                confidence_threshold=self.confidence_threshold,
+                vocab_size=self.vocab_size,
+                CL=CL,
+                ST=states.stability_threshold,
+                entropy_bound=self.entropy_bound,
+                sc_vocab_start=self.sc_vocab_start,
+                sc_vocab_end=self.sc_vocab_end,
+                tp_size=self.tp_size,
+                tp_group_name=self.tp_group_name,
+            )
+
+            # Logprobs for denoise steps that just converged (is_encoder_phase
+            # flipped False→True), stashed per tile so `scaled` is freed each tile.
+            if max_num_logprobs >= 0:
+                converged_mask = states.is_encoder_phase[tile_slots]
+                just_converged = converged_mask & ~is_committing[tile]
+                if just_converged.any():
+                    flat_logits = scaled.reshape(-1, scaled.shape[-1])
+                    argmax_tokens = scaled.argmax(dim=-1)
+                    for local_idx in just_converged.nonzero(as_tuple=True)[0]:
+                        li = local_idx.item()
+                        slot = tile_slots[local_idx]
+                        # Stash only the real canvas positions (== CL unless this
+                        # canvas was truncated near max_model_len); padded tail
+                        # positions are never emitted.
+                        k_i = int(valid_canvas_len_np[start_req + li])
+                        pos = li * CL
+                        self._pending_logprobs[slot.item()] = compute_topk_logprobs(
+                            flat_logits[pos : pos + k_i],
+                            max_num_logprobs,
+                            argmax_tokens[local_idx][:k_i],
+                        )
+
+        # Commit steps: is_committing was True at entry. Reassemble previously
+        # stashed logprobs and attach to SamplerOutput.
+        logprobs_tensors = None
+        if max_num_logprobs >= 0 and is_committing.any() and self._pending_logprobs:
+            parts_ids, parts_lp, parts_ranks = [], [], []
+            cu_gen: list[int] = []
+            flat_offset = 0
+            for i in range(num_reqs):
+                cu_gen.append(flat_offset)
+                slot = int(slots_np[i])
+                if is_decode_np[i] and slot in self._pending_logprobs:
+                    lp = self._pending_logprobs.pop(slot)
+                    parts_ids.append(lp.logprob_token_ids)
+                    parts_lp.append(lp.logprobs)
+                    parts_ranks.append(lp.selected_token_ranks)
+                    flat_offset += lp.logprobs.shape[0]
+            if parts_ids:
+                logprobs_tensors = LogprobsTensors(
+                    logprob_token_ids=torch.cat(parts_ids),
+                    logprobs=torch.cat(parts_lp),
+                    selected_token_ranks=torch.cat(parts_ranks),
+                    cu_num_generated_tokens=cu_gen,
+                )

타일링 로직은 다음과 같이 동작합니다:

1. 그룹 크기 결정: current_platform.mem_get_info()를 사용하여 현재 사용 가능한 GPU 메모리를 확인합니다. _compiled_sample_step 함수 내부에서 사용되는 임시 FP32 버퍼의 예상 크기(bytes_per_req)를 계산하고, 이를 기반으로 전체 요청 수(num_decode)를 나눌 수 있는 최대 그룹 크기(group)를 결정합니다. 이 그룹 크기는 GPU 메모리 예산의 50%를 초과하지 않도록 조정됩니다. 이 방식은 별도의 튜닝 파라미터 없이 동적으로 최적의 그룹 크기를 찾습니다.
2. 타일별 샘플링: 결정된 group 크기만큼 요청들을 묶어 for 루프를 실행합니다. 각 반복마다 해당 tile에 속하는 logits와 decode_slots 등을 _compiled_sample_step 함수에 전달합니다.
3. 로그 확률(Logprobs) 계산: 타일링된 루프 내에서, 수렴(converged)된 디노이징 단계에 대한 로그 확률을 계산하고 _pending_logprobs에 저장합니다. 이는 scaled 버퍼가 각 타일 처리 후 해제될 수 있도록 하여 메모리 사용량을 줄입니다.
4. 결과 취합: 모든 타일 처리가 완료된 후, 저장된 로그 확률들을 취합하여 최종 LogprobsTensors 객체를 생성합니다.

또한, sampled와 num_sampled 버퍼의 초기화 로직이 __call__ 메소드 시작 부분으로 이동되었습니다. 이는 타일링된 루프가 시작되기 전에 단 한 번만 수행되어, 이전 타일의 계산 결과를 다음 타일이 실수로 덮어쓰는 것을 방지합니다.

Before:

-    # Cleared inside _compiled_sample_step so prefill/non-decode slots stay 0.
+    # Clear once: the tiled loop below only scatters its own decode slots,
+    # so it must not re-clear earlier tiles' writes.
     sampled = self._sampled[:num_reqs]
     num_sampled = self._num_sampled[:num_reqs]
+    sampled.zero_()
+    num_sampled.zero_()

After:

-        # Cleared inside _compiled_sample_step so prefill/non-decode slots stay 0.
+        # Clear once: the tiled loop below only scatters its own decode slots,
+        # so it must not re-clear earlier tiles' writes.
         sampled = self._sampled[:num_reqs]
         num_sampled = self._num_sampled[:num_reqs]
+        sampled.zero_()
+        num_sampled.zero_()

이 변경은 타일링된 반복문이 각 타일의 결과만을 sampled와 num_sampled에 반영하도록 보장합니다.

왜 이게 좋은가: 성능 및 일반화 교훈

성능 향상

이 PR의 가장 큰 장점은 Diffusion-Gemma 샘플러의 임시 메모리 사용량을 획기적으로 줄였다는 점입니다. PR 설명에 따르면, H200 GPU에서 테스트했을 때 샘플러의 임시 메모리 사용량이 23.9 GiB에서 5.5 GiB로 감소했습니다. 이는 기존에는 max_num_seqs=32에서 OOM이 발생했던 상황을 해결하고, 해당 설정으로 샘플링을 가능하게 합니다.

메모리 사용량 감소는 다음과 같은 효과를 가져옵니다:

• 더 높은 동시성: 더 많은 요청을 동시에 처리할 수 있게 되어 전체 처리량(throughput)이 향상됩니다.
• 더 큰 모델 지원: 메모리 제약으로 인해 실행하지 못했던 더 큰 배치 크기나 시퀀스 길이를 가진 모델을 실행할 수 있습니다.
• 안정성 향상: OOM 오류 발생 빈도를 줄여 서비스 안정성을 높입니다.

리뷰어 masterFoad의 실험 결과에 따르면, CL=256, vocab=262144, num_decode=16 환경에서 샘플러 임시 메모리가 28,672 MiB에서 1,800 MiB로 크게 감소했습니다. 이는 타일링 기법이 매우 효과적임을 입증합니다.

일반화 교훈

이 PR은 LLM 추론 최적화에 있어 다음과 같은 중요한 교훈을 제공합니다:

1. 동적 메모리 관리의 중요성: 고정된 버퍼 크기 할당은 특정 워크로드나 하드웨어 구성에서 비효율적이거나 OOM을 유발할 수 있습니다. GPU의 가용 메모리를 실시간으로 파악하고 이에 맞춰 동적으로 리소스를 할당하는 방식(예: mem_get_info() 활용)은 시스템의 전반적인 안정성과 효율성을 크게 향상시킵니다.
2. 타일링(Tiling) 기법의 유용성: 대규모 연산을 작은 단위(타일)로 나누어 처리하는 타일링은 메모리 사용량을 제어하는 강력한 기법입니다. 특히, GPU의 임시 메모리 사용량이 워크로드 크기에 비례하여 증가하는 경우, 타일링을 통해 각 타일의 메모리 사용량을 제한함으로써 전체 메모리 사용량을 관리할 수 있습니다. 이는 _compiled_sample_step 함수가 내부적으로 여러 개의 FP32 버퍼를 유지하는 특성과 결합되어 효과를 발휘합니다.
3. 초기화 로직의 위치: 함수 내부의 상태 초기화 로직은 해당 함수의 호출 방식 및 전체 워크플로우에 미치는 영향을 신중하게 고려해야 합니다. 타일링과 같이 반복적인 호출이 발생하는 경우, 초기화 로직을 반복문 외부로 옮겨 단 한 번만 실행되도록 하는 것이 중요합니다.
4. 정확성 vs. 성능 트레이드오프: 리뷰어 masterFoad가 지적했듯이, 타일링으로 인해 확률적 샘플링 결과에서 약간의 불일치가 발생할 수 있습니다. 이는 성능 최적화 과정에서 발생할 수 있는 일반적인 현상이며, 재현성(reproducibility)이 중요한 경우 시드(seed) 관리나 결정론적 연산(deterministic operations) 사용을 고려해야 합니다. 다만, 이 PR에서는 탐욕적(greedy) 샘플링의 경우 일치함을 확인하여, 대부분의 사용 사례에서는 문제가 없을 것으로 보입니다.

리뷰 피드백 반영

리뷰어 Isotr0py의 제안에 따라, 타일 크기를 결정하는 로직이 환경 변수(env var)에 의존하는 대신 GPU의 가용 메모리를 기반으로 동적으로 결정되도록 수정되었습니다. 이는 모델별 특화된 설정을 일반적인 기능에 적용하는 것을 지양하고, 더 유연하고 자동화된 메모리 관리를 가능하게 합니다. guan404ming은 이 아이디어를 즉시 반영하여 mem_get_info()를 활용한 동적 그룹 크기 계산 로직을 구현했습니다.

결론

이번 PR은 vLLM이 Diffusion-Gemma와 같은 특정 모델 아키텍처의 메모리 병목 현상을 어떻게 효과적으로 해결하는지 잘 보여줍니다. 요청 기반 타일링과 동적 메모리 할당 기법을 통해 샘플러의 임시 메모리 사용량을 크게 줄임으로써, 더 높은 동시성과 안정성을 달성했습니다. 이는 LLM 서빙 시스템의 성능과 확장성을 개선하는 데 중요한 기여를 하며, 다른 LLM 추론 프레임워크에서도 참고할 만한 모범 사례를 제시합니다.

References

• vLLM Documentation
• Diffusion Gemma Model Card
• torch.cuda.mem_get_info
• vllm.platforms.current_platform
• vllm.v1.outputs.LogprobsTensors
• compute_topk_logprobs (This function is used internally for logprob calculation)

참고 자료

• https://pytorch.org/docs/stable/generated/torch.cuda.mem_get_info.html
• https://github.com/vllm-project/vllm/blob/main/vllm/platforms.py
• https://github.com/vllm-project/vllm/blob/main/vllm/v1/outputs.py#L15
• https://github.com/vllm-project/vllm/blob/main/vllm/sampling_ops.py#L107

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