[토론] Gemma-2B를 위한 5KB 순수 x86-64 어셈블리 엔진 (FP16, CPU에서 4.6 tok/s)
[Discussion] A 5KB pure x86-64 assembly engine for Gemma-2B (FP16, 4.6 tok/s on CPU)
핵심 요약
Gemma-2B 모델을 5KB의 순수 x86-64 어셈블리 코드로 구현하여 CPU에서 효율적으로 구동하는 프로젝트.
- 초경량 엔진 — 5.2KB의 최소한의 바이너리 크기로 Gemma-2B 모델 구동함.
- 순수 어셈블리 — C/C++ 런타임이나 PyTorch 의존성 없이 AVX2와 F16C 활용함.
- 성능 최적화 — 구형 i5 CPU에서 4.6 tok/s의 추론 속도 달성함.
- 학술적 탐구 — 현대 트랜스포머 모델을 하드웨어 수준에서 어떻게 매핑할지 연구함.
다들 안녕,
자기회귀 LLM을 돌리는 데 필요한 최소한의 베어메탈 환경을 탐구해 본 개인 프로젝트 하나 공유할게.
거대한 런타임이나 컴파일러 추상화에 의존하는 대신, Gemma-2B용 추론 엔진을 전부 순수 x86-64 어셈블리(FASM)로 직접 짰어:
-
바이너리 크기: 총 5.2 KB 플랫 머신 코드 (
gemma_engine.bin3.7 KB +mat_smp_f16c_gemm_avx2.bin1.5 KB). -
실행: 순수 AVX2 + F16C에 프리필(prefill)을 위한 커스텀 4스레드 SMP GEMM 적용. 일반 DDR4-2400 메모리에서 약 18.5 GB/s 대역폭 유지함.
-
디코딩: 구형 쿼드코어 i5 데스크탑에서 FP16 기준 4.5 ~ 4.7 토큰/s 나옴.
-
의존성: C/C++ 런타임 제로, PyTorch 제로. 파이썬 하네스는
VirtualAlloc이랑 OS 스레드 호출하려고ctypes만 썼어.
이건 llama.cpp처럼 기능 다 갖춘 툴이랑 경쟁하려는 게 아냐. 그냥 현대적인 트랜스포머 모델을 생 하드웨어에 얼마나 깔끔하게 매핑할 수 있는지, 그리고 나중에 리소스 빡빡한 마이크로컨트롤러(MCU/DSP)에서 마이크로 LLM 돌릴 때 참고 자료로 쓰려고 밑바닥부터 파본 거야.
저장소는 오픈 소스니까 확인해 봐:
-
아키텍처 노트: https://github.com/tomtsai28/PULSAR-ASM/blob/main/doc/pulsar_asm_cpu_limit_retrospective.md
코드 검토나 의견, 아니면 베어메탈 추론에 대한 생각 있으면 뭐든 환영이야.
