[vllm] vLLM Kimi-K3 MLA 디코드 성능 최적화: AITER 커널 퓨전 적용

PR 링크: vllm-project/vllm#57640 상태: Merged | 변경: +170 / -8

들어가며

vLLM의 Kimi-K3 모델은 MLA(Multi-Head Latent Attention) 구조를 사용합니다. 기존 ROCm 환경에서는 MLA 디코드 단계에서 per-layer query preparation과 KV-cache write가 각각 별도의 커널로 실행되었습니다. 구체적으로 vllm::concat_and_cache_mla_kernel과 triton_poi_fused__to_copy_cat_clamp_mul_reciprocal_view_0가 순차적으로 호출되면서 오버헤드가 발생했습니다. 본 PR은 이 두 작업을 단일 AITER 커널인 fused_qk_rope_concat_and_cache_mla로 통합하여 커널 실행 횟수를 288회에서 144회로 줄이고, 디코드 단계의 지연 시간을 약 2배 개선했습니다.

코드 분석

vllm/models/kimi_k3/amd/mla.py

핵심 변경 사항은 _fused_decode 메서드의 도입과 이를 forward 경로에 통합한 것입니다. 기존에는 두 개의 커널이 분리되어 있었으나, 이를 하나로 합치기 위해 _fused_qk_prep_supported를 통해 최적화 적용 가능 여부를 사전에 검증합니다.

Before (분리된 커널 호출):

# 기존에는 별도의 커널들이 순차적으로 실행됨
# 1. vllm::concat_and_cache_mla_kernel (KV-cache write)
# 2. triton_poi_fused_... (Q-prep/concat)

After (통합된 AITER 커널 호출):

# 단일 AITER 커널로 통합
fused_qk_rope_concat_and_cache_mla(
    ql_nope.contiguous(),
    q_pe,
    kv_c_normed,
    k_pe.squeeze(1),
    kv_cache.view(kv_cache.shape[0], -1, head_size),
    q_out,
    slot_mapping.flatten(),
    layer._k_scale,
    layer._q_scale,
    positions,
    cos_cache,
    sin_cache,
    is_neox=True,
    is_nope_first=True,
)

또한, Kimi-K3는 NoPE(No Positional Embedding) 모델이지만, AITER 커널은 RoPE를 요구합니다. 이를 해결하기 위해 cos=1, sin=0인 identity rotation을 사용하여 RoPE를 no-op으로 처리하는 영리한 방식을 채택했습니다.

왜 이게 좋은가

이번 최적화는 단순히 커널 수를 줄이는 것을 넘어, GPU 디바이스 타임을 획기적으로 단축했습니다.

• 성능 수치: MI355X (TP=8) 환경에서 MLA 레이어당 디코드 단계 지연 시간이 9.39 µs에서 4.74 µs로 약 1.98배 개선되었습니다.
• 유지보수성: 기존의 복잡한 커스텀 오프 등록 방식이나 torch.compile 가드를 제거하고, mla.py 내부에 로직을 직접 구현하여 코드 복잡도를 낮췄습니다. 또한, 다른 모델이나 백엔드에 영향을 주지 않는 'Zero Blast Radius' 설계를 통해 안전성을 확보했습니다.

일반적 교훈:

1. 커널 퓨전(Kernel Fusion): 메모리 대역폭이 병목인 연산에서 여러 커널을 하나로 합치는 것은 GPU 활용도를 높이는 가장 효과적인 방법입니다.
2. 조건부 최적화: 모든 상황에서 최적화를 적용하기보다, _fused_qk_prep_supported와 같이 제약 조건을 명확히 하여 안전한 경우에만 최적화 경로를 타도록 설계하는 것이 중요합니다.

리뷰어 피드백 반영

리뷰어 Rohan138의 제안에 따라, torch.compiler.is_compiling() 가드를 제거하고 커스텀 오프 등록 과정을 생략하여 코드를 간결하게 만들었습니다. 또한, _fused_decode 로직을 별도 파일이 아닌 mla.py 내부에 직접 통합함으로써 관리 포인트를 단일화했습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html

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