[vllm] vLLM ROCm Aiter 백엔드 성능 최적화: 불필요한 제로 필링 제거

2026년 4월 10일수정: 2026년 4월 10일

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

들어가며

대규모 언어 모델(LLM)의 추론 속도는 사용자 경험에 직접적인 영향을 미칩니다. 특히 vLLM과 같은 고성능 추론 엔진은 효율적인 GPU 활용을 통해 이를 극대화합니다. 이번 PR은 vLLM의 ROCm 환경, Aiter 백엔드에서 발생하는 불필요한 GPU 커널 실행을 제거하여 디코드(decode) 단계의 성능을 개선하는 것을 목표로 합니다.

기존 코드에서는 MLA(Multi-Layer Attention) 디코드 과정에서 출력 텐서를 초기화하기 위해 torch.zeros를 사용했습니다. 이는 GPU에 vectorized_elementwise라는 제로 필링(zero-filling) 커널을 실행하게 만듭니다. 하지만 Aiter 백엔드의 MLA 커널 파이프라인은 이 초기화된 값을 곧바로 덮어쓰기 때문에, 제로 필링 과정은 완전히 불필요하며 커널 런치 레이턴시만 증가시키는 요인이었습니다. 이 PR은 이러한 불필요한 작업을 제거하여 디코드 경로의 성능을 향상시킵니다.

코드 분석

`vllm/v1/attention/backends/mla/rocm_aiter_mla.py`

핵심 변경 사항은 rocm_aiter_mla.py 파일의 forward_mqa 함수 내에서 출력 텐서 o를 할당하는 부분입니다.

Before:

o = torch.zeros(
    B,
    self.num_heads,
    self.kv_lora_rank,
    dtype=attn_metadata.decode.attn_out_dtype,
    device=q.device,
)

After:

o = torch.empty(
    B,
    self.num_heads,
    self.kv_lora_rank,
    dtype=attn_metadata.decode.attn_out_dtype,
    device=q.device,
)

기존 코드에서는 torch.zeros를 사용하여 출력 텐서 o를 할당했습니다. torch.zeros는 텐서를 생성할 뿐만 아니라, 모든 요소를 0으로 초기화하는 추가적인 GPU 커널(vectorized_elementwise fill kernel)을 실행합니다.

하지만 Aiter MLA 커널 파이프라인(mla_decode_fwd)은 이 초기화된 o 텐서의 모든 유효한 위치를 곧바로 덮어쓰게 됩니다. 구체적으로는 다음과 같은 경우에 해당합니다:

num_kv_splits > 1일 때: Stage 2 Triton reduction 커널이 계산된 값을 o의 모든 유효한 위치에 씁니다.
num_kv_splits == 1이고 Stage 2가 실행될 때: FINAL_OUT 경로가 중간 버퍼의 내용을 o로 복사합니다.
num_kv_splits == 1이고 조기 반환(early return)이 발생할 때: logits가 o와 메모리를 공유하여 Stage 1이 직접 쓰거나, ASM 커널이 ptr_RP에 직접 씁니다.

이처럼 어떤 경우든 o 텐서의 초기값은 사용되지 않고 덮어쓰여지므로, torch.zeros를 통한 제로 필링은 명백히 불필요한 작업입니다.

수정된 코드에서는 torch.zeros 대신 torch.empty를 사용합니다. torch.empty는 단순히 텐서를 위한 메모리 공간만 할당하고, 어떠한 GPU 커널도 실행하지 않습니다. 이는 불필요한 커널 런치 레이턴시를 제거하여 디코드 경로를 최적화합니다.

이 변경은 Aiter 백엔드의 희소(sparse) MLA 구현(rocm_aiter_mla_sparse.py)에서 이미 torch.empty를 사용하고 있었기 때문에, 밀집(dense) 백엔드를 일관성 있게 만드는 효과도 있습니다.

왜 이게 좋은가

성능 향상

이 PR의 가장 큰 장점은 불필요한 GPU 커널 실행을 제거함으로써 얻는 성능 향상입니다. 구체적으로:

GPU 커널 런치 제거: 각 디코드 단계마다, 각 레이어마다 실행되던 vectorized_elementwise 제로 필링 커널 런치가 제거됩니다. GPU 커널 런치는 그 자체로 오버헤드를 가지므로, 이를 제거하는 것은 특히 짧은 커널 실행에서 큰 영향을 줄 수 있습니다.
레이턴시 감소: 불필요한 커널 실행이 사라지면서 디코드 경로의 전체 레이턴시가 감소합니다. 이는 특히 실시간 추론 서비스에서 중요한 요소입니다.

정확도 유지

PR 설명에 따르면, 이 변경은 모델의 정확도에 영향을 미치지 않습니다. Kimi-K2-Thinking 모델을 사용한 정확도 테스트 결과, torch.empty 사용 시 오히려 gsm8k 태스크에서 exact_match 정확도가 소폭 상승하는 경향을 보였습니다. 이는 제로 필링 과정에서 발생할 수 있는 미세한 부동 소수점 연산 차이나, 혹은 단순히 테스트의 무작위성에 의한 변동일 수 있습니다. 중요한 것은 성능 저하 없이 정확도가 유지된다는 점입니다.

일반적인 교훈

초기화의 비용 인지: 메모리 할당 시 torch.zeros와 torch.empty의 차이를 명확히 인지하고, 초기화가 반드시 필요한 경우가 아니라면 torch.empty를 사용하여 불필요한 연산을 피해야 합니다.
코드 경로 분석: GPU 커널이 실행되는 코드 경로를 깊이 있게 분석하여, 실제 사용되지 않는 연산이나 불필요한 초기화 과정을 찾아 제거하는 것이 성능 최적화의 핵심입니다.
백엔드 일관성: 유사한 기능을 수행하는 여러 백엔드(예: Aiter dense vs sparse) 간의 구현 일관성을 유지하는 것은 코드 관리 및 유지보수 측면에서 중요합니다.

리뷰 댓글 분석

제공된 리뷰 댓글은 주로 DCO(Developer Certificate of Origin) 서명 문제와 병합 충돌 해결에 관한 내용이었습니다. 이는 코드 자체의 기술적 타당성보다는 PR 제출 과정에서의 형식적인 부분에 대한 피드백입니다. 리뷰어인 tjtanaa가 지적한 사항을 xaguilar-amd가 신속하게 처리하여 PR이 병합될 수 있었습니다. 코드의 핵심 로직에 대한 부정적인 피드백은 없었으며, 이는 제안된 최적화가 타당했음을 시사합니다.

References

참고 자료

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

PR Analysis 의 다른글

이전글 [vllm] vLLM 성능 최적화: H2D 메모리 복사 병목 해결을 통한 추론 처리량 개선
현재글 : [vllm] vLLM ROCm Aiter 백엔드 성능 최적화: 불필요한 제로 필링 제거
다음글 [cpython] CPython JIT 최적화: 복합 마이크로 오퍼레이션(uOp)의 분해를 통한 효율성 개선

댓글

관련 포스트

PR Analysis 의 다른글