[vllm] vLLM Triton 커널 최적화: tl.constexpr 제거를 통한 JIT 컴파일 오버헤드 해결

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

들어가며

LLM 추론 엔진인 vLLM에서 merge_attn_states_kernel은 어텐션 상태를 병합하는 핵심적인 역할을 수행합니다. 하지만 기존 구현에서는 prefill_tokens_with_context 값이 tl.constexpr로 선언되어 있어, 배치마다 이 값이 달라질 때마다 Triton이 새로운 커널 변형(variant)을 JIT 컴파일해야 하는 문제가 있었습니다. 본 글에서는 이 인자를 런타임 인자로 전환함으로써 어떻게 컴파일 오버헤드를 줄이고 TTFT(Time To First Token)를 개선했는지 분석합니다.

코드 분석

Triton 커널 인자 최적화

문제의 핵심은 prefill_tokens_with_context가 커널 내부에서 단순히 마스킹 연산(token_idx < prefill_tokens_with_context)에만 사용됨에도 불구하고 tl.constexpr로 지정되어 있었다는 점입니다. Triton에서 constexpr은 컴파일 타임 상수로 취급되어, 값이 바뀔 때마다 새로운 커널이 생성됩니다.

Before

# vllm/v1/attention/ops/triton_merge_attn_states.py
    # ...
    output_scale,
    head_size,
    padded_head_size,
    output_lse is not None,
    prefill_tokens_with_context: tl.constexpr,
    # ...

After

# vllm/v1/attention/ops/triton_merge_attn_states.py
    # ...
    output_scale,
    prefill_tokens_with_context,
    head_size,
    padded_head_size,
    output_lse is not None,
    # ...

위와 같이 tl.constexpr 타입을 제거하고 일반 인자로 변경함으로써, prefill_tokens_with_context 값이 변하더라도 동일한 커널 바이너리를 재사용할 수 있게 되었습니다.

왜 이게 좋은가

성능 개선 수치

이 최적화는 단순히 코드 몇 줄을 수정한 것이지만, 실제 서비스 환경에서 극적인 성능 향상을 가져왔습니다. 벤치마크 결과에 따르면:

• 전체 평균 TTFT: 약 19.2% 개선
• p99 지연 시간: 약 11.6% 개선
• 초기 요청 시 지연 시간: 새로운 prefill_tokens_with_context 값이 들어올 때 발생하는 컴파일 스톨이 제거되어, 첫 20개 요청의 중앙값이 약 82% 감소했습니다.

일반적 교훈

Triton 개발 시 tl.constexpr은 커널 최적화(루프 언롤링, 텐서 크기 고정 등)에 필수적이지만, 남용할 경우 JIT 캐시 미스를 유발하여 오히려 성능을 저하시킵니다. 커널의 제어 흐름이나 메모리 할당에 직접적인 영향을 주지 않는 변수는 런타임 인자로 전달하는 것이 훨씬 효율적입니다. 특히 배치 단위로 값이 변하는 변수라면 반드시 런타임 인자로 분리해야 합니다.

리뷰어 피드백 반영

코드 리뷰 과정에서 prefill_tokens_with_context 인자의 위치를 다른 런타임 인자들과 함께 배치하도록 정리하여 코드 가독성을 높였습니다. 또한 불필요한 주석을 제거하여 깔끔한 PR을 유지했습니다.

참고 자료

• https://triton-lang.org/main/index.html
• https://docs.vllm.ai/en/latest/

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