[vllm] vLLM DeepSeek-OCR-2 모델의 TTFT 성능 최적화 분석

PR 링크: vllm-project/vllm#49531 상태: Merged | 변경: +32 / -36

들어가며

vLLM에서 DeepSeek-OCR-2 모델을 구동할 때, create_custom_4d_mask 함수가 CPU-bound 작업으로 인해 TTFT(Time To First Token) 지연을 유발하는 문제가 있었습니다. 기존 구현은 배치 내 각 샘플마다 루프를 돌며 CPU에서 마스크를 생성하고 GPU로 복사하는 비효율적인 방식을 취하고 있었습니다. 본 PR은 이 과정을 벡터화된 텐서 연산으로 전환하고 결과를 캐싱하여 성능을 획기적으로 개선했습니다.

코드 분석

1. vllm/model_executor/models/deepencoder2.py 최적화

기존 코드는 배치 사이즈만큼 루프를 돌며 torch.full과 인덱싱 연산을 수행했습니다. 이는 CPU에서 GPU로의 잦은 데이터 전송과 복잡한 인덱싱 오버헤드를 발생시켰습니다.

Before:

for b in range(batch_size):
    mask = torch.full((sequence_length, sequence_length), ...)
    # ... 인덱싱 및 루프 연산 ...
    masks.append(mask)
mask = torch.stack(masks, dim=0).unsqueeze(1)

After: 새로운 구현은 sequence_length에 따라 배치 불변(batch-invariant)인 마스크를 한 번만 생성합니다.

@classmethod
@lru_cache(maxsize=8)
def compute_mask_base(cls, sequence_length, dtype, device):
    # ... 벡터화된 마스크 생성 로직 ...
    return torch.where(allow, ...)[None, None]

def _create_custom_4d_mask(self, ...):
    base = self.compute_mask_base(sequence_length, dtype, device)
    return base.expand(batch_size, -1, -1, -1)

핵심은 torch.tril과 불리언 연산을 사용하여 루프 없이 마스크를 생성하고, expand를 통해 배치 차원으로 확장하는 것입니다. 이를 통해 aten::_index_put_impl_ 호출을 0으로 줄였습니다.

왜 이게 좋은가

이번 최적화는 단순히 캐싱을 도입한 것 이상의 의미가 있습니다.

1. CPU 병목 제거: torch profiler 결과에 따르면 aten::_index_put_impl_ 호출이 229,296회에서 0회로 감소했습니다. 이는 CPU에서 수행하던 복잡한 인덱싱 연산이 GPU 최적화 연산으로 대체되었음을 의미합니다.
2. 메모리 효율성: lru_cache를 사용하여 sequence_length별로 마스크를 캐싱합니다. 실제 운영 환경에서 sequence_length는 288 또는 512로 제한적이므로, 캐시 적중률이 매우 높고 메모리 점유율도 1MB 미만으로 매우 낮습니다.
3. 성능 수치: NVIDIA A10 환경에서 TTFT p50 기준 4193ms에서 2283ms로 약 1.84배(46%) 향상되었습니다.

리뷰어 피드백 반영

리뷰 과정에서 Isotr0py는 메모리 효율성을 우려했으나, LiuLi1998은 실제 시나리오에서의 제한된 sequence_length와 CPU-GPU 복사 제거의 이점을 설명하며 설득했습니다. 최종적으로 lru_cache를 클래스 메서드에 적용하여 코드를 더 간결하고 우아하게 개선했습니다.

결론

이번 PR은 모델의 특정 연산이 CPU-bound일 때, 이를 텐서 연산으로 벡터화하고 불변성을 활용해 캐싱하는 것이 얼마나 큰 성능 향상을 가져오는지 보여주는 좋은 사례입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.tril.html
• https://pytorch.org/docs/stable/generated/torch.Tensor.expand.html
• https://docs.python.org/3/library/functools.html#functools.lru_cache

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