[vllm] vLLM의 GDN 어텐션 최적화: Prefill과 Decode 배치 분리를 통한 2배 성능 향상

PR 링크: vllm-project/vllm#44700 상태: Merged | 변경: +426 / -31

들어가며

vLLM 프로젝트의 최신 PR에서는 Gated Delta Net(GDN) 어텐션 연산의 성능 병목을 해결하기 위한 중요한 최적화가 도입되었습니다. 기존에는 Prefill과 Decode가 섞인 배치(Mixed Batch)를 처리할 때, 모든 토큰을 chunk_gated_delta_rule 커널로 일괄 처리했습니다. 이 방식은 Decode 토큰 하나당 FLA_CHUNK_SIZE=64만큼의 패딩된 청크를 생성하게 하여, 실제 연산보다 훨씬 많은 낭비되는 연산(Wasted Work)을 발생시켰습니다. 본 글에서는 이 문제를 어떻게 해결했는지 코드 레벨에서 분석합니다.

코드 분석

1. Mixed Batch의 분리 전략

핵심 아이디어는 Decode 토큰을 별도의 경로로 분리하는 것입니다. GDNAttentionMetadataBuilder는 이제 Mixed 배치에서 Decode 토큰을 식별하고, 이를 fused_sigmoid_gating_delta_rule_update라는 재귀적 업데이트 커널로 라우팅합니다. 나머지 Prefill 토큰들만 기존의 chunk_gated_delta_rule을 통과하게 됩니다.

Before (기존 방식):

# 모든 토큰이 하나의 청크 커널로 처리됨
# D개의 디코드가 D개의 64-토큰 청크를 소비하여 비효율 발생
output = chunk_gated_delta_rule(mixed_qkv, ...)

After (개선된 방식):

# Decode 토큰은 재귀적 커널로, Prefill은 청크 커널로 분리
if num_decodes > 0:
    # Decode 전용 경로
    decode_out = fused_sigmoid_gating_delta_rule_update(decode_qkv, ...)
    # Prefill 전용 경로 (메타데이터 재조정)
    prefill_out = chunk_gated_delta_rule(prefill_qkv, ...)
    # 결과 병합
    return stitch_outputs(decode_out, prefill_out)

2. 메타데이터 재조정 (Rebasing)

Prefill 토큰만 청크 커널로 보내기 위해, 메타데이터 빌더는 chunk_indices와 chunk_offsets를 Prefill의 시작 위치에 맞춰 재조정(Rebase)합니다. 이를 통해 커널은 자신이 처리하는 데이터가 전체 배치의 일부임을 인지하고 정확한 연산을 수행할 수 있습니다.

왜 이게 좋은가

이번 최적화의 핵심은 '불필요한 패딩 연산 제거'입니다.

• 성능 수치: 커널 마이크로벤치마크(B200 기준)에서 기존 4361 µs가 2256 µs로 단축되어 약 1.93배의 속도 향상을 보였습니다.
• 실제 처리량: 엔드투엔드 테스트에서 초당 토큰 처리량(Total tok/s)이 131,976에서 164,156으로 +24.4% 증가했습니다.

일반적 교훈

1. 커널 특성에 맞는 데이터 라우팅: 모든 데이터를 하나의 '범용' 커널에 넣는 것이 항상 최선은 아닙니다. 연산 특성(재귀적 vs 청크 기반)에 따라 데이터를 분리하는 것이 훨씬 효율적일 수 있습니다.
2. 메타데이터의 유연성: 복잡한 연산에서 성능을 높이려면, 커널에 전달되는 메타데이터(인덱스, 오프셋 등)를 동적으로 재구성하는 능력이 필수적입니다.

리뷰 과정에서 contiguous() 호출이 불필요하다는 피드백이 있었고, 이를 즉시 반영하여 불필요한 메모리 복사 오버헤드까지 제거하는 등 코드 품질을 높였습니다.

참고 자료

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

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