[vllm] vLLM 성능 최적화: Sparse MLA 인덱스 변환 공유를 통한 3.9배 성능 향상

PR 링크: vllm-project/vllm#59464 상태: Merged | 변경: +147 / -9

들어가며

최신 거대 언어 모델(LLM) 아키텍처인 GLM-5.3은 효율적인 추론을 위해 Sparse MLA(Multi-Head Latent Attention) 기법을 사용합니다. 하지만 기존 구현에서는 각 레이어마다 동일한 논리적 인덱스(logical indices)를 물리적 KV 슬롯으로 변환하는 연산이 반복적으로 수행되고 있었습니다. 이는 모델의 레이어 수가 많아질수록 불필요한 연산 오버헤드를 발생시키는 병목 지점이었습니다. 본 PR은 이 중복된 변환 과정을 레이어 그룹 단위로 공유함으로써 추론 성능을 획기적으로 개선했습니다.

코드 분석

1. vllm/v1/attention/backends/mla/flashattn_mla_sparse.py

핵심 변경 사항은 forward_mqa 함수 내에서 인덱스 변환 로직을 조건부로 분기 처리한 것입니다. 기존에는 모든 레이어에서 triton_convert_req_index_to_global_index를 호출했지만, 이제는 index_group이 존재하고 레이어 그룹 내에서 재사용이 가능한 경우 최적화된 경로를 사용합니다.

Before:

topk_indices, valid_counts = triton_convert_req_index_to_global_index(
    attn_metadata.req_id_per_token[:num_actual_toks],
    attn_metadata.block_table,
    topk_indices,
    BLOCK_SIZE=attn_metadata.block_size,
    BLOCK_STRIDE_ROWS=block_stride_rows,
    NUM_TOPK_TOKENS=topk_indices.shape[1],
    return_valid_counts=True,
)

After:

if (
    index_group is not None
    and index_group.num_layers > 1
    and attn_metadata.num_decode_tokens == num_actual_toks
):
    topk_indices, valid_counts = self._convert_logical_to_physical_topk(
        topk_indices,
        attn_metadata,
        block_stride_rows=block_stride_rows,
        return_valid_counts=True,
    )
else:
    # 기존 로직 유지

2. tests/v1/attention/test_sparse_mla_backends.py

최적화된 로직이 기존의 독립적인 변환 결과와 동일한지 검증하기 위해 test_flashattn_shared_indices_match_independent_conversion 테스트 케이스가 추가되었습니다. MagicMock을 사용하여 triton_convert_req_index_to_global_index의 호출 횟수를 추적함으로써, 실제로 인덱스 변환이 공유되고 있는지 확인합니다.

왜 이게 좋은가

이번 최적화는 GLM-5.3 모델의 메인 모델 포워드 패스에서 인덱스 변환 횟수를 78회에서 21회로 대폭 줄였습니다. 벤치마크 결과에 따르면, 토큰 수에 따라 다르지만 약 3.5배에서 3.9배에 달하는 커널 성능 향상을 보여줍니다.

일반적 교훈:

1. 중복 연산 제거: 동일한 입력에 대해 반복되는 GPU 커널 호출은 오버헤드가 큽니다. 이를 캐싱하거나 공유하는 구조로 변경하는 것만으로도 큰 성능 이득을 얻을 수 있습니다.
2. 데이터 재사용: 메모리 접근 패턴이 동일하다면, 연산 결과를 재사용하는 것이 GPU 메모리 대역폭과 연산 자원을 아끼는 핵심 전략입니다.
3. 검증의 중요성: 성능 최적화 시 로직의 정확성을 보장하기 위해 Mock 객체를 활용한 호출 횟수 검증과 기존 결과값과의 비교 테스트는 필수적입니다.

참고 자료

• vLLM GitHub Repository
• Triton Documentation

참고 자료

• https://github.com/vllm-project/vllm
• https://triton-lang.org/main/index.html

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