[vllm] vLLM DeepSeek-V4 성능 최적화: 불필요한 torch.full 커널 제거로 1.88배 속도 향상

PR 링크: vllm-project/vllm#50298 상태: Merged | 변경: +44 / -23

들어가며

vLLM의 DeepSeek-V4 모델 구현에서 성능 병목을 유발하던 불필요한 GPU 커널 호출을 제거한 최적화 사례를 소개합니다. 특히 combine_topk_swa_indices 연산 과정에서 매번 새로운 텐서를 생성하기 위해 호출되던 torch.full 커널이 성능 저하의 주범이었으며, 이를 사전 할당된 out 텐서를 재사용하는 방식으로 개선하여 1.88배의 성능 향상을 달성했습니다.

코드 분석

1. vllm/models/deepseek_v4/common/ops/cache_utils.py

기존 코드에서는 combine_topk_swa_indices 함수가 호출될 때마다 torch.full을 사용하여 새로운 텐서를 할당했습니다. 이는 GPU 메모리 할당 비용과 커널 실행 오버헤드를 발생시켰습니다.

Before:

combined_indices = torch.full(
    (num_tokens, combined_topk),
    fill_value=-1,
    dtype=torch.int32,
    device=topk_indices.device,
)
combined_lens = torch.empty(
    num_tokens, dtype=torch.int32, device=topk_indices.device
)

After: out 파라미터를 추가하여 호출자가 미리 할당된 텐서를 전달할 수 있도록 변경했습니다.

if out is None:
    # ... 기존 로직 ...
else:
    combined_indices, combined_lens = out

2. vllm/models/deepseek_v4/nvidia/flashmla.py

_forward_prefill 함수에서 workspace_manager를 통해 미리 할당된 메모리 버퍼를 가져와 combine_topk_swa_indices에 전달하도록 수정했습니다.

After:

workspace = workspace_manager.get_simultaneous(
    ((chunk_size, chunk_M, q.shape[-1]), torch.bfloat16),
    ((self.max_num_batched_tokens, combined_topk), torch.int32),
    ((self.max_num_batched_tokens,), torch.int32),
)
kv, combined_indices_out, combined_lens_out = workspace
# ...
combined_indices, combined_lens = combine_topk_swa_indices(
    ...,
    out=(combined_indices_out, combined_lens_out),
)

왜 이게 좋은가

이 최적화의 핵심은 GPU 메모리 할당의 최소화와 커널 호출 오버헤드 제거입니다. torch.full은 GPU에서 메모리를 할당하고 초기화하는 커널을 실행하는데, 이는 특히 작은 연산 단위가 빈번하게 호출되는 prefill 단계에서 큰 비용이 됩니다.

성능 측정 결과:

• GPU Latency: 0.093 ms → 0.049 ms (1.88배 속도 향상)
• Wall Latency: 0.094 ms → 0.050 ms (1.86배 속도 향상)

교훈

1. Pre-allocation: 반복적으로 호출되는 연산에서는 텐서를 매번 생성하지 말고, workspace를 통해 미리 할당된 버퍼를 재사용하십시오.
2. Kernel Overhead: GPU 연산에서 커널 실행 횟수 자체를 줄이는 것이 전체 파이프라인의 레이턴시를 낮추는 가장 효과적인 방법 중 하나입니다.
3. API Design: 라이브러리 함수를 설계할 때 out 파라미터를 제공하여 호출자가 메모리 관리 전략을 선택할 수 있게 하는 것이 좋습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.full.html
• https://pytorch.org/docs/stable/generated/torch.empty.html

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