메모리 장벽을 깨다: Kog Inference Engine으로 3,000 토큰/초 달성

현재 자율 AI 에이전트의 병목 현상은 모델의 지능뿐만 아니라 반복 속도에 있습니다. 검사, 계획, 편집, 테스트의 순차 루프를 수행하는 에이전트에게 토큰 생성에 소요되는 시간이 바로 루프 속도를 결정합니다. 워크플로우에 50,000 토큰이 필요할 때, 100 토큰/초와 3,000 토큰/초의 차이는 8분 대기와 20초 응답의 차이와 같습니다.

Kog AI는 Kog Inference Engine (KIE)의 기술 미리보기를 공개했으며, 표준 데이터센터 GPU가 이전에 전용 추론 하드웨어에만 허용되던 속도를 달성할 수 있음을 보여줍니다. 8× AMD MI300X GPU에서 요청당 3,000 출력 토큰/초, 8× NVIDIA H200에서 2,100 토큰/초(2B 모델을 FP16으로 사용) 를 달성함으로써, Kog는 총 처리량에서 단일 요청 디코드 속도로 초점을 전환하고 있습니다.

메모리 대역폭 병목 현상

저배치 자동회귀 디코딩에서는 주요 제약이 계산 능력(FLOPS)이 아니라 메모리 대역폭입니다. 생성되는 각 토큰마다 모델의 활성 가중치는 High Bandwidth Memory (HBM)에서 연산 프로세서로 이동해야 합니다.

현대 GPU는 피크 FLOPS와 메모리 대역폭 사이에 큰 불균형을 가지고 있습니다. 예를 들어, NVIDIA H200은 HBM 대역폭당 약 400 FLOP를 제공합니다. 배치 크기 1에서 토큰 생성은 연산 강도가 매우 낮아(~1 FLOP/byte in FP16) 시스템은 GPU의 계산 능력을 소진하기 전에 메모리 대역폭 한계에 도달합니다.

Kog의 전략은 전체 8-GPU 서버 노드를 단일 연속 메모리 스트리밍 머신으로 취급하여, 8개의 카드 전체 HBM 대역폭을 집계해 초당 생성되는 토큰 수를 최대화하는 것입니다.

"마이크로초 손실" 제거

3,000 토큰/초에 도달하려면 토큰당 예산은 단 333 마이크로초입니다. 표준 추론 스택(vLLM 또는 TensorRT-LLM 등)에서는 이 예산이 시스템 오버헤드에 소모됩니다:

  • Kernel Boundaries: 커널을 시작하고 정리하는 과정 및 스케줄러 라운드트립은 작업당 몇 마이크로초를 추가할 수 있습니다. 25개 레이어에 걸치면 속도가 크게 제한됩니다.
  • CPU Scheduling: 호스트 측 로직 및 GPU-CPU 통신이 실행 지연을 초래합니다.
  • Grid Synchronization: 정규화나 샘플링을 위한 전역 GPU 동기화가 소중한 마이크로초를 소모합니다.
  • Inter-GPU Collectives: 표준 텐서 병렬 처리에서는 레이어당 여러 AllReduce 연산이 필요하며, 이는 느릴 수 있습니다.

Kog는 엔진, GPU 코드, 모델 아키텍처를 공동 설계하여 단일 지연 최적화 파이프라인으로 이러한 문제를 해결합니다.

핵심 기술 혁신

  1. The Monokernel Runtime: 커널 시퀀스 대신 Kog는 전체 디코드 경로에 대해 단일 지속 GPU 프로그램을 사용합니다. 이는 커널 경계를 없애고 호스트 측 스케줄링을 제거하여 시스템이 중단 없이 가중치를 스트리밍할 수 있게 합니다.
  2. KCCL (Kog Collective Communication Layer): 어셈블리 수준에서 튜닝된 맞춤형 통신 레이어입니다. 공급업체 라이브러리가 컬렉티브에 8µs를 사용할 수 있는 반면, KCCL은 3µs 이하로 유지하며 모노커널 스케줄에 원활히 통합됩니다.
  3. Laneformer Architecture: 지연 텐서 병렬성(DTP)을 특징으로 하는 모델 설계입니다. DTP는 디코딩의 의존 구조를 변경하여 디바이스 간 통신이 계산과 겹치게 하며, 중요한 경로를 차단하지 않게 합니다.
  4. Topology-Aware Memory Access: AMD MI300X와 같이 컴퓨트 다이(XCDs)와 I/O 다이(IODs) 사이에 비균일한 접근 경로가 있는 하드웨어에서, Kog는 버퍼를 특정 HBM 위치에 매핑하여 지연과 왜곡을 최소화합니다.

최첨단 MoE 모델으로 확장

기술 미리보기가 2B 모델을 사용하지만, 핵심 논지는 활성 파라미터 바이트 이동량에 기반하며 전체 파라미터 수와는 무관합니다. 이는 토큰당 전체 파라미터 중 일부만 활성화되는 Mixture-of-Experts (MoE) 모델에 특히 강력한 접근법입니다.

Kog는 이 스택을 더 큰 오픈 웨이트 MoE에 적용하고 FP8/FP4 양자화를 활용하면 표준 데이터센터 GPU에서 최첨단 모델을 1,000–5,000 토큰/초 범위로 끌어올릴 수 있다고 추정합니다. 서드파티 모델의 아키텍처를 변경할 수 없더라도(즉 DTP를 사용할 수 없더라도) 모노커널과 KCCL의 조합으로 통신을 가중치 스트리밍과 겹칠 수 있습니다.

커뮤니티 관점 및 비판

고성능 주장과 마찬가지로, 이번 발표는 기술 커뮤니티 내에서 논쟁을 일으켰습니다. 일부 비평가들은 2B 모델에서 속도를 보여주는 것이 최첨단 모델(30B+ 파라미터) 성능의 대리 지표가 아니며, 확장은 선형 예측보다 더 어려울 수 있다고 주장합니다.

"2B 파라미터 모델에 대한 이러한 주장은 1코어에서 4코어까지의 선형 확장성을 본 뒤 256코어가 256배 속도를 제공할 것이라고 가정하는 것과 비슷합니다... 더 큰 모델으로 확장하는 것이 예상보다 더 어려울 것이라는 느낌이 듭니다."

다른 사람들은 이러한 극단적인 속도의 유용성을 의문시하며, 토큰 품질이 생성 속도보다 중요하거나 400-500 토큰/초가 대부분 인간 중심 코딩 작업에 이미 충분하다고 지적합니다. 그러나 "추론 루프"에서 작동하는 자율 에이전트에게 생산성 최전선은 지능 × 반복 속도에 의해 정의되며, 회복의 매 마이크로초가 가치 있습니다.

Sources