본문으로 건너뛰기

[vllm] [vLLM] ROCm 환경에서의 DeepSeek-V2/V3 성능 극대화를 위한 MLA 최적화 분석

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

들어가며

DeepSeek-V2 및 V3 모델은 MLA(Multi-head Latent Attention) 구조를 통해 추론 효율성을 극대화합니다. 하지만 이러한 복잡한 어텐션 구조를 AMD GPU(ROCm) 환경에서 최적으로 실행하기 위해서는 하드웨어 특성에 맞는 커널 튜닝과 데이터 레이아웃 최적화가 필수적입니다.

이번 PR([ROCm][Deepseek] dsv3.2 further optimization)은 vLLM의 ROCm 백엔드에서 DeepSeek 모델의 추론 성능을 한 단계 더 끌어올리기 위한 최적화들을 담고 있습니다. 주요 변경 사항으로는 KV 캐시 인덱서 레이아웃 최적화, ROCm 전용 paged_mqa_logits 통합, 그리고 FP8 Sparse MLA 지원이 포함되었습니다.

코드 분석: 핵심 변경 사항

1. DeepSeek-V2 모델의 RoPE 연산 최적화

기본적으로 RoPE(Rotary Positional Embedding) 연산 시 split과 cat 연산이 빈번하게 발생하면 메모리 대역폭에 병목이 생길 수 있습니다. 이번 변경에서는 ROCm 환경에서 불필요한 텐서 분할 및 재결합을 줄이는 방향으로 코드가 수정되었습니다.

Before:

q_pe, q_nope = torch.split(
    q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
)
# ... (중략) ...
q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
# ... (중략) ...
q = torch.cat([q_pe, q_nope], dim=-1)

After:

if current_platform.is_rocm():
    # ... (중략) ...
    rotary_emb(
        positions, q[..., : self.rope_dim], k[..., : self.rope_dim].unsqueeze(1)
    )
else:
    # 기존 CUDA 로직 유지

왜 좋은가? ROCm 경로에서는 q와 k 텐서를 직접 슬라이싱하여 rotary_emb에 전달함으로써, 명시적인 torch.split과 torch.cat 호출로 인한 오버헤드를 방지합니다. 이는 특히 컴파일러가 최적화하기 어려운 메모리 복사 연산을 줄여줍니다.

2. Triton 기반의 Sparse MLA 인덱스 변환 커널 도입

Sparse MLA를 지원하기 위해 요청(Request) 인덱스를 글로벌 KV 캐시 인덱스로 변환하는 전용 Triton 커널이 추가되었습니다. 이는 Paged Attention 환경에서 흩어져 있는 블록들을 효율적으로 참조하기 위함입니다.

핵심 코드 (vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py):

@triton.jit
def _convert_req_index_to_global_index_kernel(
    req_id_ptr, block_table_ptr, token_indices_ptr, cu_seqlens_ptr, out_ptr,
    max_num_blocks_per_req: tl.constexpr, BLOCK_SIZE: tl.constexpr, BLOCK_N: tl.constexpr,
    bt_stride0, bt_stride1, ti_stride0, ti_stride1,
):
    # ... (중략) ...
    block_id = tok // BLOCK_SIZE
    inblock_off = tok % BLOCK_SIZE

    bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
    base = tl.load(bt_ptr, mask=valid_block, other=0)

    out_val = tl.where(
        is_invalid_tok | (~valid_block), 0, base * BLOCK_SIZE + inblock_off
    )
    # ... (중략) ...

이 커널은 각 토큰의 인덱스를 계산할 때 block_table을 참조하여 물리적 메모리 주소로 매핑합니다. tl.where를 사용하여 유효하지 않은 토큰(-1 등)을 안전하게 처리하며, 병렬 처리를 통해 인덱싱 속도를 극대화했습니다.

3. 인덱서 블록 사이즈 유연화

기존에는 ROCm에서 블록 사이즈가 1로 고정되어 있었으나, 이번 PR을 통해 64 사이즈를 함께 지원하도록 확장되었습니다.

변경 사항 (vllm/v1/attention/backends/mla/indexer.py):

@staticmethod
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
    return [1, 64] if current_platform.is_rocm() else [64]

리뷰어 @tjtanaa의 질문에 대해 개발자는 "ROCm Sparse MLA에서 커널은 블록 사이즈에 크게 구애받지 않지만, NVIDIA와의 호환성 및 성능을 위해 64를 지원하는 것이 유리하다"고 답변했습니다. 이는 다양한 하드웨어 구성에서 최적의 성능을 낼 수 있는 유연성을 제공합니다.

왜 이게 좋은가? (성능 및 교훈)

  1. FP8 Sparse MLA 지원: DeepSeek-V3와 같은 최신 모델은 FP8 정밀도를 적극 활용합니다. 이번 PR은 ROCM_AITER_MLA_SPARSE 백엔드에서 FP8 지원 범위를 넓혀 메모리 대역폭 절감과 연산 속도 향상을 동시에 달성했습니다.
  2. 런타임 오버헤드 감소: 메타데이터 빌드 과정을 런타임이 아닌 metadata_builder 단계로 옮김으로써, 매 스텝 발생하는 CPU-GPU 동기화 및 연산 부하를 줄였습니다.
  3. 안정성 확보: 특정 Head Size(예: 32)에서 발생하던 Illegal Memory Access 문제를 명시적인 assert와 조건 처리를 통해 해결했습니다.

테스트 결과: GSM8K 벤치마크에서 exact_match 기준 0.9477의 높은 점수를 기록하며, 최적화 이후에도 모델의 정확도가 안정적으로 유지됨을 확인했습니다.

결론

이번 PR은 ROCm 환경에서 DeepSeek 모델을 운영하는 엔지니어들에게 매우 중요한 업데이트입니다. 특히 Triton을 활용한 커스텀 커널 최적화와 메모리 레이아웃 개선은 대규모 언어 모델 추론 시스템에서 성능을 결정짓는 핵심 요소입니다. vLLM의 이러한 행보는 AMD 하드웨어 생태계에서도 최신 SOTA 모델들을 효율적으로 서빙할 수 있는 기반이 됩니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글