Hybrid JIT-CUDA Graph Optimization for Low-Latency Large Language Model Inference

[기술 요약] Hybrid JIT–CUDA Graph를 이용한 저지연 LLM 추론 최적화

1. 문제 정의 (Problem Statement)

대규모 언어 모델(LLM)의 실시간 배포, 특히 대화형 AI와 같은 인터랙티브 환경에서는 **추론 지연 시간(Inference Latency)**과 **지연 시간 변동성(Latency Variance)**을 줄이는 것이 핵심 과제입니다. 현재 LLM 추론은 다음과 같은 기술적 한계에 직면해 있습니다:

• 커널 런칭 오버헤드 (Kernel Launch Overhead): 자기회귀(Autoregressive) 디코딩 과정에서 매 토큰 생성 시마다 수많은 GPU 커널을 반복적으로 호출해야 하며, 이 과정에서 CPU-GPU 간의 통신 및 호스트 측의 디스패치(Dispatch) 오버헤드가 누적됩니다.
• 정적 vs 동적 실행의 트레이드오프:
  • CUDA Graph: 커널 실행 경로를 미리 캡처하여 재사용함으로써 CPU 오버헤드를 획기적으로 줄일 수 있지만, 텐서 크기가 고정되어야 하고 제어 흐름이 결정론적(Deterministic)이어야 한다는 제약이 있어 가변적인 LLM의 특성(가변 길이 어텐션, 확률적 샘플링 등)을 수용하기 어렵습니다.
  • JIT(Just-In-Time) 컴파일: 동적인 제어 흐름과 가변적인 텐서 형태를 처리할 수 있지만, 여전히 매 실행 시마다 호스트 측의 스케줄링 및 디스패치 오버헤드가 발생합니다.

2. 제안 방법론 (Methodology: Hybrid Runtime Framework)

본 논문은 정적 실행의 효율성과 동적 실행의 유연성을 결합한 하이브리드 JIT-CUDA Graph 런타임을 제안합니다.

• 연산의 이원화 (Partitioning): 트랜스포머 추론 파이프라인을 두 가지 영역으로 분리합니다.
  • 정적 영역 (Static Domain): 선형 투영(Linear Projections), 레이어 정규화(Layer Norm), 행렬 곱셈(MatMul) 등 텐서 형태가 일정하고 결정론적인 연산은 CUDA Graph로 캡처하여 재사용합니다.
  • 동적 영역 (Dynamic Domain): 토큰 샘플링(Sampling), KV-캐시 업데이트, 가변 길이 어텐션 등 데이터에 따라 제어 흐름이 변하는 연산은 JIT 컴파일을 통해 처리합니다.
• 비동기 그래프 생성 (Asynchronous Graph Generation): 추론 실행을 중단하지 않고 백그라운드에서 새로운 시퀀스 길이에 대한 CUDA Graph를 캡처하는 메커니즘을 도입합니다. 이를 통해 새로운 길이에 대한 그래프 생성 비용을 실제 추론 과정과 오버랩(Overlap)시켜 지연을 최소화합니다.
• 비동기 제어 구조 (Asynchronous Controller): 별도의 CUDA 스트림을 사용하여 'Context Generator(JIT 기반 동적 연산)'와 'Graph Generator(CUDA Graph 기반 정적 연산)'를 독립적으로 운영하며, IPC(Inter-Process Communication)를 통해 데이터를 주고받습니다.

3. 주요 기여 (Key Contributions)

1. 하이브리드 아키텍처 설계: LLM 추론의 정적/동적 요소를 분리하여 CUDA Graph의 성능 이점과 JIT의 유연성을 동시에 확보했습니다.
2. 비동기 캡처-재생 메커니즘: 추론 중에도 백그라운드에서 그래프를 생성하여, 새로운 시퀀스 길이에 대응할 때 발생하는 오버헤드를 분산시켰습니다.
3. 지연 시간 및 변동성 감소: 짧은 시퀀스(Interactive workload) 환경에서 TTFT(Time-to-First-Token)와 P99(Tail Latency)를 효과적으로 개선했습니다.

4. 실험 결과 (Results)

NVIDIA H100 GPU 환경에서 LLaMA-2 7B 모델을 대상으로 실험한 결과는 다음과 같습니다:

• TTFT (Time-to-First-Token) 개선: PyTorch Eager 방식 대비 최대 66.0%, TensorRT-LLM 대비 최대 5.42배 빠른 성능을 보였습니다. 특히 프롬프트 길이가 길어질수록 TensorRT-LLM보다 훨씬 안정적인 선형적 스케일링을 보여주었습니다.
• P99 Tail Latency 감소: 토큰 생성 시 발생하는 꼬리 지연 시간(Tail Latency)을 TensorRT-LLM 대비 평균 20.2%, PyTorch Eager 대비 43.5% 감소시켜, 사용자 경험 측면에서 훨씬 예측 가능한(Deterministic) 응답 속도를 제공했습니다.
• Ablation Study: 비동기 그래프 재생을 끄면 TTFT가 17.5% 증가하고, JIT를 제거하면 샘플링 로직의 오버헤드로 인해 지연 시간이 28% 증가함을 확인하여 두 기술의 상호 보완적 역할을 입증했습니다.

5. 의의 및 결론 (Significance & Conclusion)

본 연구는 기존의 최적화 방식(TensorRT-LLM 등)이 가진 '정적 그래프 중심'의 한계를 극복하고, **실제 사용자가 체감하는 인터랙티브 LLM 환경(짧은 문장 생성, 빠른 응답 필요)**에 최적화된 새로운 실행 모델을 제시했습니다.

비록 메모리 사용량(그래프 개수에 따른 증가)과 단일 GPU 범위라는 한계가 있으나, 향후 **Shape Bucketing(텐서 크기 그룹화)**이나 멀티 GPU 확장을 통해 더욱 범용적인 저지연 LLM 추론 엔진으로 발전할 수 있는 가능성을 열었습니다.