[vllm] [vLLM] GPU-CPU 동기화 병목 제거: prepare_chunk_indices 최적화 분석

PR 링크: vllm-project/vllm#38361 상태: Merged | 변경: +None / -None

들어가며

고성능 LLM 추론 엔진인 vLLM에서 GDN(Gated Delta Net) 어텐션 메커니즘의 Prefill 단계는 매우 중요한 성능 지점입니다. 최근 vLLM 프로젝트에 반영된 한 PR은 GDN Prefill 과정에서 미세하지만 치명적인 병목을 유발하던 GPU-CPU 동기화(Synchronization) 문제를 해결했습니다.

문제의 핵심은 prepare_chunk_indices 함수 내부에 숨어있던 .tolist() 호출이었습니다. PyTorch에서 GPU 텐서에 .tolist()를 호출하면, 해당 데이터를 CPU로 가져오기 위해 GPU의 모든 작업이 완료될 때까지 CPU가 기다리는 '블로킹(Blocking)' 현상이 발생합니다. 비록 @tensor_cache를 통해 캐싱을 시도했더라도, 매 스텝의 첫 번째 호출에서 발생하는 이 동기화는 파이프라인 전체를 멈추게(Stall) 만들기에 충분했습니다.

본 글에서는 이 병목을 어떻게 식별하고, 아키텍처적으로 어떻게 개선했는지 상세히 살펴보겠습니다.

코드 분석: 무엇이 문제였는가?

1. 기존의 병목 지점 (Before)

기존 코드에서는 FLA(Fused Linear Attention) 연산이 실행될 때마다 내부에서 인덱스를 준비했습니다.

# vllm/model_executor/layers/fla/ops/chunk_delta_h.py (Before)
def chunk_gated_delta_rule_fwd_h(
    ...,
    cu_seqlens: torch.Tensor | None = None,
):
    # 내부에서 매번 호출됨
    chunk_indices = (
        prepare_chunk_indices(cu_seqlens, chunk_size)
        if cu_seqlens is not None
        else None
    )
    # ... 이후 연산 수행

여기서 호출되는 prepare_chunk_indices는 내부적으로 GPU 텐서인 cu_seqlens를 기반으로 계산을 수행한 뒤, 최종 결과를 CPU 리스트로 변환하기 위해 .tolist()를 호출합니다. Nsight Systems 프로파일링 결과, 이 지점에서 DtoH(Device to Host) 메모리 복사와 함께 CPU가 대기하는 현상이 명확히 관찰되었습니다.

2. 메타데이터 기반 사전 계산 (After)

이번 개선의 핵심은 **"어차피 CPU에 이미 있는 정보라면, GPU 연산 중에 CPU로 다시 가져오지 말고 미리 계산해서 넘겨주자"**는 것입니다.

먼저, GDNAttentionMetadataBuilder에서 cu_seqlens_cpu를 활용해 인덱스를 미리 계산합니다.

# vllm/v1/attention/backends/gdn_attn.py (After)
# Metadata 생성 시점에 CPU에서 미리 계산
chunk_indices, chunk_offsets = None, None
if cu_seqlens is not None:
    # cu_seqlens_cpu는 이미 CPU에 있으므로 동기화가 발생하지 않음
    chunk_indices = prepare_chunk_indices(cu_seqlens_cpu, FLA_CHUNK_SIZE)
    chunk_offsets = prepare_chunk_offsets(cu_seqlens_cpu, FLA_CHUNK_SIZE)

return GDNAttentionMetadata(
    ...,
    chunk_indices=chunk_indices,
    chunk_offsets=chunk_offsets,
)

그 후, FLA 연산 체인의 모든 함수가 이 미리 계산된 인덱스를 선택적으로 전달받을 수 있도록 수정되었습니다.

# vllm/model_executor/layers/fla/ops/chunk_delta_h.py (After)
def chunk_gated_delta_rule_fwd_h(
    ...,
    chunk_indices: torch.Tensor | None = None, # 외부에서 주입 가능
    chunk_offsets: torch.Tensor | None = None,
):
    # 외부에서 주입받았다면 prepare_chunk_indices 호출(및 .tolist() 동기화)을 건너뜀
    if chunk_indices is None and cu_seqlens is not None:
        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
    
    # ...

왜 이게 좋은 최적화인가?

1. GPU-CPU 동기화의 완전한 제거

Nsight Systems 프로파일링 결과에 따르면, 수정 전에는 첫 번째 GDN 레이어가 실행될 때 CPU 타임이 튀는 현상이 있었으나, 수정 후에는 모든 GDN 블록의 CPU 점유 시간이 균일해졌습니다. 특히 **DtoH memcpy가 0%**가 되었다는 점은 불필요한 데이터 이동이 완전히 사라졌음을 의미합니다.

2. 하드코딩 제거 및 상수화

기존 코드 곳곳에 흩어져 있던 chunk_size=64라는 매직 넘버를 FLA_CHUNK_SIZE라는 상수로 통합했습니다. 이는 코드의 가독성을 높일 뿐만 아니라, 향후 하드웨어 특성에 맞춰 청크 크기를 조정해야 할 때 유지보수성을 크게 향상시킵니다.

3. 하위 호환성 유지

chunk_indices를 Optional 파라미터로 설계하여, 이 최적화가 적용되지 않은 다른 모델(KDA, OLMo Hybrid 등)에서는 기존의 @tensor_cache 로직을 그대로 사용할 수 있도록 배려했습니다. 이는 대규모 프로젝트에서 기존 기능을 깨뜨리지 않으면서 성능을 개선하는 모범적인 방식입니다.

성능 결과

• Prefill Forward Time: 약 410ms에서 385ms로 단축되었습니다.
• Throughput & Accuracy: 대규모 벤치마크(Qwen3.5-397B) 결과, 정확도 저하 없이 성능 수치가 안정적으로 유지됨을 확인했습니다.

결론 및 교훈

딥러닝 커널 최적화에서 가장 흔히 간과하는 것이 바로 Host-Device 간의 암시적 동기화입니다. tolist(), item(), nonzero()와 같은 함수들은 편리하지만, 루프 내부나 빈번하게 호출되는 함수 내에 있을 경우 전체 시스템의 병목이 됩니다.

이번 PR은 "데이터가 어디에 있는가?"를 정확히 파악하고, 이미 Host(CPU)에 있는 데이터를 재활용함으로써 불필요한 동기화를 제거한 훌륭한 사례입니다. 시니어 엔지니어라면 라이브러리가 제공하는 추상화 너머에서 실제로 어떤 하드웨어 상호작용이 일어나는지 항상 주시해야 함을 다시 한번 일깨워줍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.tolist.html
• https://triton-lang.org/main/index.html
• https://developer.nvidia.com/nsight-systems

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