본문으로 건너뛰기

[vllm] vLLM, DeepSeek-V3.2 모델의 ROCm 성능 최적화: CPU 측 마이크로 최적화 3가지 분석

PR 링크: vllm-project/vllm#42982 상태: Merged | 변경: +59 / -29

들어가며

최근 vLLM 프로젝트의 GitHub 저장소에는 DeepSeek-V3.2 모델을 ROCm 환경에서 실행할 때 CPU 측의 성능을 개선하는 PR이 올라왔습니다. 이 PR은 GPU 커널 수정 없이, 디스패치 로직 개선을 통해 3가지 독립적인 마이크로 최적화를 적용했습니다. 특히 5k/500/64 워크로드에서 스텝당 3.10%의 성능 향상을 달성하며, 이는 대규모 언어 모델(LLM)의 추론 속도 향상에 중요한 기여를 합니다. 본 글에서는 이 PR의 변경 사항을 상세히 분석하고, 각 최적화가 왜 효과적인지, 그리고 어떤 교훈을 얻을 수 있는지 살펴보겠습니다.

코드 분석

이번 PR은 크게 세 가지 주요 최적화를 포함하며, 각각 vllm/model_executor/models/deepseek_v2.pyvllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py 파일에서 이루어졌습니다.

1. Router bf16 dispatch 최적화 (vllm/model_executor/models/deepseek_v2.py)

DeepSeek-V3.2 모델의 MoE(Mixture of Experts) 라우팅 경로에서 불필요한 bfloat16tofloat32_copy_kernel_cuda 커널 호출을 제거했습니다. 기존에는 AITER의 biased_grouped_topk 커널이 bf16을 직접 처리할 수 있음에도 불구하고, vLLM이 fp32 형태의 라우터 로짓을 요구하여 bf16 -> fp32 캐스트가 발생했습니다. 이 PR에서는 게이트(gate)의 출력 데이터 타입을 명시적으로 가중치(weight)의 데이터 타입과 일치시켜 이 불필요한 캐스트를 제거했습니다.

Before:

# Before
self.gate.set_out_dtype(torch.float32)

After:

# After
self.gate.set_out_dtype(self.gate.weight.dtype)   # bf16

이 변경으로 인해 ROCm AITER MoE 라우팅 경로에서 스텝당 약 58회의 bfloat16tofloat32_copy_kernel_cuda 커널 실행이 제거되었습니다. 이는 GPU 커널 실행 횟수를 줄여 직접적인 성능 향상을 가져옵니다.

2. Sparse-MLA 메타데이터 캐싱 (vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py)

AttentionMetadataBuilder.build() 호출 시마다 실행되는 aiter::get_mla_metadata_v1 커널은 (num_tokens, max_query_len, num_heads, min(seq_lens, topk_tokens))와 같은 입력값에 따라 결정되는 워크로드 스케줄링을 계산합니다. 이 PR에서는 이러한 입력값들의 해시(fingerprint)를 CPU 측에서 계산하고, 동일한 입력에 대해서는 커널 실행을 건너뛰도록 캐싱 로직을 추가했습니다.

변경 전 (기존 로직):

매번 get_mla_metadata_v1 호출

변경 후 (캐싱 적용):

# ... (이전 코드) ...

        metadata_key = (
            num_tokens,
            int(common_attn_metadata.max_query_len),
            self._num_attention_heads,
            clamped_seq_lens.tobytes(),
        )
        if metadata_key != self._prev_metadata_key:
            from aiter import get_mla_metadata_v1

            get_mla_metadata_v1(
                # ... (get_mla_metadata_v1 호출 인자들) ...
            )
            self._prev_metadata_key = metadata_key

# ... (이후 코드) ...

이 최적화는 특히 긴 컨텍스트 디코딩(seq_lens >= topk_tokens=2048) 시 캐시 히트율이 약 100%에 달하여, 스텝당 get_mla_metadata_v1 커널 호출 횟수를 1023회에서 64회로 약 94% 감소시켰습니다. 이는 반복적인 계산을 피함으로써 상당한 성능 향상을 이끌어냅니다.

3. Sparse-MLA tail-fill 최적화 (vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py)

기존에는 req_id_per_token_buffer, paged_kv_indices, paged_kv_indptr와 같은 버퍼들을 전체적으로 fill_(0) 하는 작업이 수행되었습니다. 이 PR에서는 이러한 전체 버퍼 채우기 대신, 실제로 변경된 부분(shrink-tail)만 채우도록 로직을 수정했습니다. 특히 paged_kv_indptr.fill_(0)는 이후 cumsum + scalar_broadcast 연산에 의해 완전히 덮어쓰여지므로, 해당 fill_(0) 호출을 완전히 제거했습니다.

Before:

# Before (기존 로직의 일부)
self.req_id_per_token_buffer.fill_(0)
self.paged_kv_indices.fill_(0)
self.paged_kv_indptr.fill_(0)

After:

# After (변경 후 로직의 일부)
# Only re-zero the shrink-tail. paged_kv_indptr is fully rewritten
# by the cumsum below. paged_kv_indices entries past new_indices_extent
# are never read (the attention kernel only touches the ranges
# defined by paged_kv_indptr).
new_req_extent = int(req_id_per_token.shape[0])
new_indices_extent = num_tokens * self.topk_tokens
if self._prev_req_extent > new_req_extent:
    self.req_id_per_token_buffer[new_req_extent : self._prev_req_extent].fill_(
        0
    )
if self._prev_indices_extent > new_indices_extent:
    self.paged_kv_indices[new_indices_extent : self._prev_indices_extent].fill_(
        0
    )
self._prev_req_extent = new_req_extent
self._prev_indices_extent = new_indices_extent
self.req_id_per_token_buffer[:new_req_extent].copy_(
    torch.from_numpy(req_id_per_token), non_blocking=True
)

이 변경은 불필요한 메모리 연산을 줄여 성능을 개선합니다. 특히 paged_kv_indptr의 경우, 완전히 새로 계산되는 부분이므로 초기화(fill_(0)) 자체가 불필요했습니다.

왜 이게 좋은가?

이 PR에서 적용된 최적화들은 다음과 같은 이유로 매우 효과적입니다:

  1. CPU 측 최적화: GPU 커널 수정 없이 CPU 측의 디스패치 및 데이터 준비 로직을 개선함으로써, ROCm 환경에서의 성능 병목 현상을 완화했습니다. 이는 GPU 커널 최적화만큼이나 중요하며, 때로는 더 쉽게 적용 가능합니다.
  2. 불필요한 연산 제거: 불필요한 데이터 타입 캐스트(bf16 -> fp32) 및 메모리 버퍼 초기화(fill_(0))를 제거하여 연산량을 줄였습니다. 이는 직접적인 성능 향상으로 이어집니다.
  3. 캐싱 전략 도입: 반복적으로 동일한 계산이 발생하는 경우, CPU 측에서 이를 감지하고 건너뛰는 캐싱 메커니즘을 도입했습니다. 이는 특히 get_mla_metadata_v1과 같이 입력에 따라 결정되는 계산에서 큰 효과를 발휘합니다.
  4. 측정된 성능 향상: 5k/500/64 워크로드에서 스텝당 평균 1.288ms, 즉 3.10%의 성능 향상을 달성했습니다. 이는 956 스텝 동안 총 1.232초의 시간 단축으로 이어집니다.
  5. 정확도 유지: GSM8K 데이터셋을 이용한 정확도 검증 결과, flexible-extract 정확도가 0.9378 ± 0.0067로 이전과 동일하며, 어떠한 정확도 회귀도 발생하지 않았습니다. 이는 성능 개선과 함께 모델의 신뢰성을 보장합니다.

일반적 교훈:

  • LLM 추론 성능 최적화는 GPU 커널뿐만 아니라 CPU 측의 데이터 준비, 디스패치 로직에서도 큰 잠재력을 가집니다.
  • 불필요한 데이터 변환 및 메모리 연산을 식별하고 제거하는 것은 항상 좋은 성능 개선 전략입니다.
  • 반복적인 계산 패턴을 파악하고 캐싱을 적용하는 것은 성능을 크게 향상시킬 수 있습니다.
  • 성능 최적화 시에는 반드시 정확도 회귀 여부를 철저히 검증해야 합니다.

리뷰 피드백 반영

리뷰 과정에서 몇 가지 중요한 피드백이 있었습니다:

  • 정확도 우선: tjtanaa 리뷰어는 성능 향상보다 정확도 보장을 우선시하며, 특히 num_shot 20으로 GSM8K 데이터셋 전체에 대한 정확도 재검증을 요청했습니다. 이는 성능 최적화가 모델의 근본적인 성능을 해치지 않도록 하는 중요한 원칙입니다. 이후 frida-andersson이 해당 테스트를 수행하여 정확도 회귀가 없음을 확인했습니다.
  • 코드 간결성: AndreasKaratzas 리뷰어는 불필요한 주석을 제거하고 코드의 가독성을 높일 것을 제안했습니다. 특히 ROCm 환경에서 torch.zeros 초기화가 일반적이라는 점, 변수명이 충분히 설명적이라는 점 등을 지적하며 코드의 간결성을 강조했습니다. 이에 따라 fadadfc 커밋에서 상당수의 주석이 제거되었습니다.
  • 초기화 방식: paged_kv_indptr 버퍼의 초기화 방식에 대한 논의가 있었습니다. frida-anderssontorch.empty 대신 torch.zeros를 사용하고, build 함수 내에서 첫 호출 시 분기 처리 대신 할당 시점에 zeros를 사용하는 것이 더 안전하고 간결하다고 제안했으며, 이는 후속 커밋에서 반영되었습니다.

이러한 리뷰 피드백은 PR의 완성도를 높이는 데 크게 기여했습니다.

References

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글