AI 소프트웨어 스택 개요

개요

AI 소프트웨어 스택(AI Software Stack)은 머신러닝 모델을 하드웨어 가속기에서 효율적으로 실행하기 위한 계층형 소프트웨어 시스템을 총칭한다. 최상위의 프레임워크(TensorFlow, PyTorch)부터 최하위의 하드웨어 드라이버까지, 여러 계층이 유기적으로 협력하여 개발자가 하드웨어 세부 사항을 몰라도 높은 성능을 달성할 수 있게 한다.

현대 AI 모델은 수십억에서 수조 개의 파라미터를 가지며, 대규모 행렬 연산이 필수적이다. 이러한 모델을 효율적으로 실행하려면 연산 융합(Operator Fusion), 메모리 최적화, 병렬화, 하드웨어 특화 코드 생성 등 복잡한 최적화가 필요하다. AI 소프트웨어 스택은 이러한 최적화를 자동화하여 GPU, TPU, NPU 등 다양한 가속기에서 최적의 성능을 발휘할 수 있게 한다.

동일한 모델이라도 학습, 오프라인 배치 추론, 온라인 저지연 추론은 병목이 다르기 때문에 스택의 무게중심도 달라진다. 학습에서는 통신 라이브러리와 자동 미분 경로가 중요하고, 서비스 추론에서는 그래프 캡처, 정밀도 변환, KV cache 관리, 배포 런타임이 더 중요하다. 따라서 AI 소프트웨어 스택을 볼 때는 단순히 "프레임워크 아래에 라이브러리가 있다"는 계층도뿐 아니라, 어떤 계층이 어떤 워크로드 병목을 담당하는지 함께 봐야 한다.

핵심 개념

AI 소프트웨어 스택 계층 구조

AI 소프트웨어 스택은 크게 네 가지 계층으로 구성된다:

서빙 엔진과 컴파일러의 경계

실무에서는 프레임워크 아래에 "컴파일러/런타임"만 있는 것이 아니라, 그 위나 옆에 배치되는 서빙 엔진 계층도 중요하다. 예를 들어 TensorRT-LLM, Triton Inference Server, vLLM 같은 엔진은 모델 배치 정책, KV cache 관리, 요청 스케줄링, 동시성 제어를 담당하며, 내부적으로는 다시 TensorRT, CUDA, NCCL, cuDNN 같은 하위 계층을 호출한다.

NVIDIA CUDA 생태계

CUDA(Compute Unified Device Architecture)는 NVIDIA가 개발한 병렬 컴퓨팅 플랫폼으로, AI 소프트웨어 스택의 가장 광범위한 생태계를 구축하고 있다.

CUDA 스택 구성:

CUDA 주요 라이브러리:

CUDA 프로그래밍 모델:
- Grid-Block-Thread: 3단위 병렬화 구조
- Shared Memory: 블록 내 공유 메모리, 용량과 분할 정책은 아키텍처별로 다름
- Constant/Texture Memory: 읽기 전용 캐시 메모리
- Warp Scheduling: 32개 스레드 단위 실행 (SIMT)

AMD ROCm 생태계

ROCm(Radeon Open Compute)는 AMD가 개발한 오픈소스 GPU 컴퓨팅 플랫폼으로, NVIDIA CUDA와 호환성을 유지하면서도 독자적인 발전을 거듭하고 있다.

ROCm 스택 구성:

ROCm 주요 라이브러리:

HIP(Heterogeneous-compute Interface for Portability):
- CUDA 코드를 ROCm에서 실행하기 위한 포팅 레이어
- hipify 도구로 CUDA → HIP 자동 변환 가능
- HIP API는 이식 계층이고, 실제 최적화 수준은 ROCm 라이브러리와 컴파일러 품질에 크게 좌우됨

Google TPU 소프트웨어 스택

TPU 소프트웨어 스택은 Google의 JAX/XLA 중심 생태계로, TPU 하드웨어에 최적화된 컴파일러와 런타임을 제공한다.

TPU 스택 구성:

XLA의 주요 최적화:
- 연산 융합: 여러 연산을 하나의 커널로 결합
- 메모리 최적화: 중간 결과물의 불필요한 할당 제거
- 하드웨어 특화: TPU systolic array에 최적화된 코드 생성
- JIT 컴파일: 실행 시점에서 동적 최적화

JAX(JAX Advanced eXtensions):
- NumPy와 유사한 API로 함수형 프로그래밍 지원
- jit: XLA를 통한 JIT 컴파일
- vmap: 벡터화된 맵핑 (배치 처리 자동화)
- pmap: 다중 TPU/GPU 병렬 처리
- grad: 자동 미분 (모든 연산에서)

비교/분석

주요 AI 소프트웨어 스택 비교

대표 실행 계층 기능 비교

프로그래밍 모델 비교

동작 원리

AI 추론 파이프라인 소프트웨어 흐름

1. 모델 로드: 프레임워크가 저장된 모델 가중치를 로드
2. 그래프 최적화: 컴파일러가 연산 그래프를 분석하고 최적화
3. 커널 생성: 하드웨어 특화 커널을 동적 또는 사전 생성
4. 메모리 할당: GPU/TPU 메모리에 텐서 할당
5. 커널 실행: 드라이버를 통해 하드웨어에 커널 전달 및 실행
6. 결과 반환: 출력 텐서를 호스트 메모리로 복사

연산 융합 동작 원리

연산 융합은 여러 독립적인 연산을 하나의 커널로 결합하여 메모리 접근을 줄이는 핵심 최적화 기법이다:

융합 예시 (Conv + Bias + ReLU):
- 융합 전: Conv 커널 실행 → 메모리 쓰기 → Bias 커널 실행 → 메모리 쓰기 → ReLU 커널 실행
- 융합 후: 하나의 커널에서 Conv → Bias → ReLU 순차 실행 → 메모리 쓰기 1회

이를 통해 메모리 대역폭 소비를 최대 60%까지 줄이고, 커널 시작 오버헤드를 제거하여 2-3배 성능 향상을 달성할 수 있다.

실제 서비스 환경에서는 여기에 정적 shape 고정, KV cache 레이아웃 정규화, collective overlap, 연산 그래프 캡처 같은 추가 최적화가 더해진다. 즉, 컴파일러가 커널 수를 줄이는 역할을 하고, 런타임과 서빙 엔진이 요청 단위의 스케줄링 비용을 줄이는 역할을 나눠 담당한다.

하드웨어 가상화 및 리소스 관리

현대 AI 소프트웨어 스택은 여러 워크로드가 동일한 하드웨어를 공유할 수 있는 가상화 기능을 제공한다:

TPU 계열에서는 MIG처럼 사용자가 직접 하드웨어를 세밀하게 분할한다기보다, PJRT와 클러스터 런타임이 디바이스 메시에 작업을 배치하는 방식이 중심이다. ROCm 역시 MIG와 동일한 의미의 표준 분할 계층보다는 프로세스 격리, 컨테이너 배치, 통신 토폴로지 최적화가 운영 관점에서 더 중요하다.

장단점

NVIDIA CUDA 생태계

장점:
- 가장 성숙하고 광범위한 생태계 (15년 이상)
- cuDNN, TensorRT 등 고성능 라이브러리
- 풍부한 문서, 커뮤니티, 도구 지원
- 멀티 GPU/노드 확장에 최적화 (NCCL)
- 프로덕션 배포에 최적 (TensorRT, Triton)

단점:
- NVIDIA GPU 전용 (하드웨어 종속)
- 완전한 오픈소스가 아님 (cuDNN, NCCL 비공개)
- 라이선스 비용 발생
- 경쟁사 하드웨어로의 이전 어려움

AMD ROCm 생태계

장점:
- 완전한 오픈소스 플랫폼
- NVIDIA CUDA와의 높은 호환성 (HIP)
- AMD Instinct GPU의 높은 메모리 대역폭 활용
- 라이선스 비용 없음
- AMD GPU의 높은 가성비

단점:
- NVIDIA 대비 생태계 성숙도 부족
- 일부 라이브러리 성능 격차 (MIOpen vs cuDNN)
- 드라이버 안정성 이슈 보고
- 프레임워크 호환성 문제 간헐적 발생

Google TPU/JAX 생태계

장점:
- TPU 하드웨어에 최적화된 소프트웨어 스택
- JAX의 함수형 프로그래밍 모델 (우아한 병렬화)
- XLA의 강력한 그래프 최적화
- Pod 수준 대규모 확장 지원
- Google Cloud 인프라와의 긴밀한 통합

단점:
- TPU 전용 (하드웨어 종속)
- JAX 학습 곡선이 높음
- 로컬 개발 환경 제한
- 커스텀 커널 작성 어려움
- TPU 외 하드웨어 지원 제한

관련 기술

참고 표준 및 사양

주요 프레임워크와의 통합

관련 문서

• 텐서 코어 아키텍처
• TPU 상세
• CUDA 상세
• ROCm 상세
• TPU 소프트웨어 스택
• AI 컴파일러 개요

주요 참고 문헌

1. NVIDIA. "CUDA C++ Programming Guide." NVIDIA Developer Documentation.
2. NVIDIA. "CUDA Toolkit Documentation." NVIDIA Developer Zone.
3. AMD. "ROCm Documentation." AMD ROCm Docs.
4. AMD. "HIP Programming Guide." AMD Developer Documentation.
5. Google. "TPU Architecture." Google Cloud Documentation.
6. Google. "JAX Documentation." Google JAX GitHub.
7. Google. "XLA: Accelerated Linear Algebra." OpenXLA Project.
8. NVIDIA. "cuDNN Library User Guide." NVIDIA Developer Zone.
9. NVIDIA. "TensorRT Developer Guide." NVIDIA Developer Zone.
10. AMD. "MIOpen Library." AMD ROCm Documentation.

핵심 정리

1. AI 소프트웨어 스택은 프레임워크, 서빙 엔진, 컴파일러/런타임, 라이브러리, 드라이버가 연동되는 구조이며, 개발자가 하드웨어 세부 사항을 몰라도 높은 성능을 달성할 수 있게 한다.

2. NVIDIA CUDA는 가장 성숙한 생태계를 구축하고 있으며, cuDNN, TensorRT, NCCL 등 고성능 라이브러리와 풍부한 도구 지원으로 프로덕션 배포의 표준으로 자리잡았다.

3. AMD ROCm은 완전한 오픈소스 플랫폼으로, HIP를 통해 CUDA 코드의 높은 호환성을 제공하면서도 AMD GPU의 높은 메모리 대역폭을 활용할 수 있다.

4. Google TPU/JAX 생태계는 함수형 프로그래밍 모델과 XLA 컴파일러의 강력한 최적화를 통해 TPU 하드웨어에서 최적의 성능을 발휘하며, Pod 수준의 대규모 확장을 지원한다.

5. 연산 융합, 메모리 최적화, 하드웨어 특화 코드 생성 등 AI 컴파일러의 핵심 최적화 기법들이 모든 스택에서 공통적으로 적용되며, 스택별로 고유한 최적화 전략을 가지고 있다.