Stanford CS336 Language Modeling from Scratch: Inference Engines and Full-Stack Innovation

Stanford CS336 Language Modeling from Scratch: Inference Engines and Full-Stack Innovation

Inference: The Engine Turning Electricity into Intelligence

Inference는 훈련된 언어 모델을 정적인 수학 객체(연산의 Directed Acyclic Graph)에서 실제 도구로 변환하는 핵심 단계입니다. 기본적인 inference 엔진과 GPU 커널을 이해하는 것은 "full‑stack innovation"에 필수적인데, 이 구성 요소들이 모델의 실제 효율성, 지연 시간, 그리고 프로덕션에서의 기능을 결정하기 때문입니다.

The Lifetime of a Token

요청이 inference 시스템에 들어오면 다음과 같은 파이프라인을 따릅니다:

  1. Scheduling: 요청이 특정 GPU로 라우팅되며, 경우에 따라 prefill 작업과 decode 작업을 서로 다른 머신에 분산시킬 수 있습니다.
  2. KV Cache Lookup: 시스템은 이전에 동일한 요청이나 그 변형이 있었는지 확인해 계산을 절약합니다.
  3. Execution: 핵심 머신러닝 코드가 실행되며, 모델 크기에 따라 노드 혹은 GPU 간에 병렬화됩니다.
  4. Token Generation: 시스템이 토큰을 출력하고, 이후 stop sequence와 safety check을 위해 처리됩니다.

Prefill vs. Decode Workloads

Inference는 근본적으로 서로 다른 계산 특성을 가진 두 단계로 구성됩니다:

  • Prefill: 초기 입력 프롬프트(예: 10,000 토큰)를 처리해 초기 활성화를 계산합니다. 이는 compute‑bound이며, 역전파가 없는 학습과 유사하게 높은 FLOPs를 사용합니다.
  • Decode: 토큰을 하나씩 생성합니다. 전체 모델을 메모리에서 불러와 단일 토큰을 만들기 때문에 memory bandwidth‑bound이며, 요구되는 FLOPs는 상대적으로 적습니다.

이 차이 때문에 최신 스택은 종종 prefill과 decode를 서로 다른 워커 집합이나 특화된 하드웨어(예: prefill에 NVIDIA GPU, decode에 LPU/Grok 칩)로 분리합니다.

Production Challenges and System Optimizations

하루에 수조 개의 토큰을 서비스하면 작은 규모에서는 드러나지 않던 미묘한 버그와 병목 현상이 나타납니다. 여기에는 커널 오류로 인해 모델이 토큰을 무한히 반복하는 "doom loops" 혹은 초기화되지 않은 GPU 메모리를 읽는 커널의 off‑by‑one 오류로 인한 예상치 못한 언어 전환(예: 영어에서 중국어) 등이 포함됩니다.

KV Cache Management

처리량을 최대화하려면 Key‑Value(KV) 캐시를 효율적으로 관리해야 합니다. GPU 메모리는 제한적이므로 계층형 저장 방식을 사용합니다:

  • GPU Memory: 가장 빠른 접근 속도, 가장 제한된 공간.
  • CPU DRAM: 느리지만 덜 자주 사용되는 활성화를 오프로드하는 데 사용.
  • SSD/Disk: 가장 느리며 세션 데이터의 장기 저장에 활용.

이는 고전적인 운영체제 메모리 관리와 유사하며, 보통 Least Recently Used(LRU) 히어스틱을 적용해 어떤 활성화를 느린 저장소로 내보낼지 결정합니다.

Cache‑Aware Disaggregation

간단하지만 효과적인 최적화는 캐시 히트율에 따라 요청을 라우팅하는 것입니다. "신선한" 요청(캐시 히트율 낮고 prefill 비용 높음)은 한 세트의 GPU에, "따뜻한" 요청(캐시 히트율 높고 prefill 비용 낮음)은 다른 세트에 보내면 서비스 속도가 최대 40%까지 향상될 수 있습니다.

Megakernels: Achieving "Speed of Light" Inference

전통적인 inference 엔진은 연산을 하나씩 실행합니다(예: Norm 커널 → MatMul 커널). 이는 커널 실행 오버헤드와 "tail effects"(GPU가 배치 내 가장 긴 시퀀스가 끝날 때까지 다음 연산을 시작하지 못함) 때문에 큰 지연을 초래합니다.

The Megakernel Approach

Megakernel은 여러 연산을 하나의 커널로 합쳐 GPU를 순차 연산자가 아닌 거대한 분산 시스템처럼 활용합니다. 이를 통해 다음과 같은 연산 겹침이 가능합니다:

  • KV Cache를 로드하면서 QKV 프로젝션과 RoPE 스케일링이 아직 진행 중인 경우.
  • 어텐션 연산이 완전히 끝나기 전에 O 프로젝션 가중치를 로드.

저수준 CUDA 제어를 위한 ThunderKittens 라이브러리를 사용하면 Megakernel은 H100 GPU에서 최대 72% 대역폭 활용률을 달성하며 거의 "빛의 속도"에 가까운 성능을 보여줍니다.

Parcae: Stabilized Recurrent Architectures

대부분의 현재 LLM은 파라미터와 데이터를 늘려 스케일링하지만, Parcae 연구는 재귀(recursion)(트랜스포머 블록을 루프하는 방식)로 스케일링하는 방법을 탐구합니다.

The Stability Problem

단순 루프 트랜스포머는 매우 불안정합니다. 학습률 같은 하이퍼파라미터를 약간만 바꿔도 손실 급등과 NaN이 발생합니다. 이 불안정성은 변환 행렬의 스펙트럴 반경 때문이며, 행렬 노름이 1보다 크면 루프를 통해 활성화가 기하급수적으로 폭발합니다.

The Parcae Solution

Parcae는 재귀 시스템의 A와 B 행렬을 재파라미터화하여 학습을 안정화합니다:

  • Matrix A: 시스템이 안정적이고 활성화가 폭발이 아니라 점차 감소하도록 음의 대각 행렬로 제한.
  • Matrix B: 간단한 선형 노름으로 제어.

이 수학적 제약은 안정적인 손실 곡선을 보장하고, 전통적인 트랜스포머보다 파라미터당 더 높은 품질을 달성하게 합니다.

Scaling Laws for Recurrence

초기 연구 결과에 따르면 재귀는 데이터와 파라미터와 함께 스케일링되어야 합니다. 즉, 학습 데이터 양이 증가함에 따라 재귀 횟수도 늘려야 계산 효율성을 유지할 수 있습니다. 이는 고정된 파라미터 예산 하에서 루프 사전학습을 수행하면 더 높은 품질의 모델을 얻을 수 있음을 시사합니다.

SUMMARY: Guest lecturer Dan Fu discusses the critical role of inference engines and GPU kernels in transforming LLMs from mathematical objects into usable intelligence, introducing optimizations like Megakernels and the Parcae recurrent architecture.

TITLE: Stanford CS336 Language Modeling from Scratch: Inference Engines and Full-Stack Innovation

Sources