BitNet(1.58비트)을 위한 제로 의존성 C 추론 엔진 구축 - Xeon CPU에서 36 tok/s 달성하며 얻은 교훈
Building a zero-dependency C inference engine for BitNet (1.58-bit) - lessons from hitting 36 tok/s on a Xeon CPU
핵심 요약
의존성 없는 순수 C99로 BitNet 1.58비트 모델을 CPU에서 효율적으로 구동하는 추론 엔진 개발기입니다.
- 제로 의존성 — Python, CUDA, BLAS 없이 GCC와 make만으로 빌드됨
- 네이티브 3진법 SIMD — AVX2/AVX-512를 활용해 정수 레지스터에서 직접 연산
- 메모리 대역폭 한계 — 연산 최적화보다 DRAM 대역폭이 추론 속도의 병목 현상으로 작용
- 성능 최적화 — C11 원자 연산을 사용해 스레드 동기화 오버헤드를 거의 제로로 구현
지난 몇 달 동안 저는 순수 C99(Python, CUDA, BLAS 없이 오직 GCC와 make만 사용)로 CPU 우선 추론 엔진을 밑바닥부터 구축해왔습니다. 핵심은 무거운 런타임 오버헤드 없이 1.58비트 3진법 모델을 네이티브로 실행하는 것이었습니다.
현재 Intel Xeon에서 4개의 스레드를 사용하여 BitNet b1.58-2B-4T 모델로 36.25 tok/s를 기록하고 있습니다.
빌드 과정에서 얻은 몇 가지 기술적 세부 사항입니다:
- 네이티브 3진법 SIMD: BitNet 가중치는 바이트당 4개씩 패킹되어 있습니다(값은 -1, 0, +1). 수학 연산 전에 float32로 언패킹하는 대신, VNNI 명령어(vpdpbusds)를 사용하여 정수 레지스터에 직접 누적하는 커스텀 AVX2 및 AVX-512 루틴을 사용합니다.
- 최소한의 런타임 오버헤드: 스레드 풀은 무거운 뮤텍스 경합 대신 C11 원자 연산과 spin-then-yield 백오프를 사용하여 토큰 생성 중 스레드 동기화 오버헤드가 사실상 제로에 가깝습니다.
- 제로 의존성: OpenAI 호환 API 엔드포인트를 직접 제공하는 단일 독립형 바이너리로 컴파일됩니다.
지금까지 얻은 가장 큰 교훈은 DRAM의 한계입니다. 저수준 행렬 곱셈 커널을 최적화하는 데 몇 주를 보냈지만, 배치 사이즈 1에서의 디코드 속도는 메모리 대역폭에 묶여 있습니다. Xeon 테스트 박스에서 이론상 메모리 대역폭의 약 95%를 사용하고 있기 때문에, 여러 시퀀스에 걸쳐 배치 처리를 하지 않는 한 더 빠른 연산 커널이 엔드투엔드 토큰 지연 시간을 단축시키지는 못합니다.
레포: https://github.com/shifulegend/project-zero
다른 CPU 아키텍처(특히 AMD Zen이나 ARM NEON)에서 어떤 토큰 속도가 나오는지, 혹은 로컬 3진법 추론에서 메모리 대역폭 한계를 어떻게 다루고 있는지 궁금합니다.
