[vllm] vLLM 성능 최적화: token_to_req_indices 캐싱을 통한 6배 성능 향상

PR 링크: vllm-project/vllm#47474 상태: Merged | 변경: +34 / -25

들어가며

vLLM은 고성능 추론 엔진으로서 효율적인 메모리 관리와 커널 연산이 핵심입니다. 최근 DeepSeek-V4 모델을 지원하는 과정에서, token_to_req_indices를 생성하는 로직이 여러 컴포넌트(Sparse MLA, SWA, Compressor)에서 독립적으로 수행되고 있음을 발견했습니다. 이는 불필요한 CPU-GPU 간 데이터 복사를 유발하여 병목 현상을 일으키고 있었습니다. 본 PR은 이 매핑 로직을 CommonAttentionMetadata로 중앙화하고 캐싱함으로써, 커널 수준에서 5~6배의 성능 향상을 달성했습니다.

코드 분석

1. CommonAttentionMetadata의 중앙화 (vllm/v1/attention/backend.py)

가장 핵심적인 변경은 매핑 로직을 CommonAttentionMetadata 클래스 내부로 옮긴 것입니다. 기존에는 각 모듈이 개별적으로 np.repeat와 copy_를 수행했으나, 이제는 캐시를 확인하고 존재하지 않을 때만 연산을 수행합니다.

# After: 중앙화된 캐싱 로직
def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor:
    if self._token_to_req_indices_cache is not None:
        return self._token_to_req_indices_cache[:num_tokens]
    
    # ... (생성 및 복사 로직) ...
    self._token_to_req_indices_cache = buffer[: max(num_mapped_tokens, num_tokens)]
    return self._token_to_req_indices_cache[:num_tokens]

2. 각 모듈의 리팩토링 (sparse_mla.py, compressor.py, sparse_swa.py)

기존의 복잡한 NumPy 연산과 pin_memory 호출이 제거되고, 단순히 common_attn_metadata의 메서드를 호출하는 형태로 간소화되었습니다.

# Before (sparse_mla.py)
- starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)
- seg_lengths = np.diff(starts)
- req_id_per_token = np.repeat(...)
- self.req_id_per_token_buffer.copy_(...)

# After
+ req_id_per_token = cm.token_to_req_indices(self.req_id_per_token_buffer)

왜 이게 좋은가

이번 최적화는 '중복 연산 제거'와 '데이터 전송 최소화'라는 두 가지 측면에서 매우 효과적입니다.

1. 성능 수치: 벤치마크 결과, decode bs32 환경에서 기존 194.97us 소요되던 작업이 32.73us로 단축되어 약 5.96배의 속도 향상을 보였습니다. 특히 prefill 단계에서도 일관되게 5~6배의 성능 개선이 확인되었습니다.
2. 교훈:
   • 데이터 지역성(Data Locality): GPU 연산에서 CPU-GPU 간 데이터 전송은 매우 비싼 작업입니다. 동일한 데이터를 여러 번 전송하는 대신, 한 번 전송한 후 재사용하는 전략이 필수적입니다.
   • 추상화의 힘: 메타데이터 관리 클래스에 연산을 위임함으로써, 개별 모델 구현체는 비즈니스 로직에만 집중할 수 있게 되었고 코드 유지보수성도 크게 향상되었습니다.

리뷰 과정에서 CUDA Graph와의 호환성을 위해 메모리 주소를 고정하는 작업이 중요하게 다루어졌으며, 결과적으로 안정성과 성능을 모두 확보하는 방향으로 최적화가 완료되었습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.copy_.html
• https://docs.vllm.ai/en/latest/

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