[vllm] vLLM의 GLM5.2 성능 최적화: Triton 커널 융합을 통한 E2E Throughput 향상

PR 링크: vllm-project/vllm#46862 상태: Merged | 변경: +167 / -0

들어가며

최신 대규모 언어 모델(LLM) 추론 엔진인 vLLM에서 GLM5.2 모델의 성능을 극대화하기 위한 최적화가 진행되었습니다. 기존의 파이프라인은 Q RoPE -> cat -> FP8 quant -> q_scale fold 순서로 연산이 분리되어 있어, 메모리 접근 횟수가 많고 커널 호출 오버헤드가 존재했습니다. 본 PR은 이 연산들을 하나의 Triton 커널로 융합하여 메모리 효율을 높이고 추론 처리량(Throughput)을 1.9%에서 3.3%까지 향상시켰습니다.

코드 분석

1. vllm/model_executor/layers/sparse_attn_indexer.py

핵심 변경 사항은 _fused_indexer_q_rope_quant_kernel의 도입입니다. 기존에는 여러 단계로 나뉘어 있던 연산을 triton.jit을 사용하여 단일 커널로 통합했습니다.

Before (개념적 흐름):

# 개별 커널 호출
q_rope = apply_rope(q)
q_fp8 = quantize_fp8(q_rope)
weights_out = fold_scale(weights, q_scale)

After (Triton Kernel):

@triton.jit
def _fused_indexer_q_rope_quant_kernel(...):
    # RoPE 연산 수행
    r0 = (x0 * cos - x1 * sin).to(tl.bfloat16).to(tl.float32)
    # FP8 양자화 및 스케일 계산
    q_scale = tl.math.exp2(tl.ceil(tl.log2(scale_raw)))
    # 가중치 스케일 폴딩 및 저장
    tl.store(weights_out, weight * q_scale * softmax_scale * head_scale)

이 융합을 통해 중간 결과물을 GPU Global Memory에 쓰고 다시 읽는 과정을 생략하고, 레지스터 수준에서 연산을 처리하여 메모리 대역폭 병목을 해결했습니다.

2. vllm/model_executor/models/deepseek_v2.py

모델의 forward 패스에서 새로운 융합 커널을 사용하도록 로직을 분기했습니다.

# After: 조건부 융합 커널 호출
elif self.use_fused_indexer_q and q.dtype == torch.bfloat16:
    q_fp8, weights = fused_indexer_q_rope_quant(
        positions, q, rotary_emb.cos_sin_cache, weights, ...
    )
    return self.indexer_op(hidden_states, q_fp8, k, weights)

왜 이게 좋은가

성능 개선

벤치마크 결과, 기존 대비 Output token throughput이 약 60.01 tok/s에서 61.17 tok/s로 향상되었으며, 전체적인 E2E 성능이 유의미하게 개선되었습니다. 특히 메모리 접근이 잦은 RoPE와 양자화 연산을 융합함으로써, GPU의 연산 유닛 활용도를 극대화했습니다.

기술적 교훈

1. Kernel Fusion의 중요성: 메모리 집약적인 연산(Memory-bound)들은 커널을 융합하여 데이터 이동을 최소화하는 것이 성능 최적화의 핵심입니다.
2. Triton의 활용: CUDA C++을 직접 작성하는 것보다 Triton을 사용하면 유지보수성과 생산성을 유지하면서도 컴파일러 최적화를 통해 높은 성능을 얻을 수 있습니다.
3. 리뷰 피드백 반영: 1e-4와 1e-10 사이의 epsilon 값 논의에서 보듯, 기존 시스템과의 일관성(Consistency)을 유지하는 것이 모델의 정확도(Accuracy)를 보장하는 데 필수적입니다.

참고 자료

• Triton Documentation — Triton JIT 컴파일러 공식 문서
• vLLM GitHub Repository — vLLM 프로젝트 메인 페이지

참고 자료

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

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