Rust로 구현한 1비트 LLM 추론 엔진: CPU에서 150+ TPS 달성

섹션별 상세

1비트 양자화 모델을 CPU에서 구동하기 위해 Rust 기반의 제로 의존성 추론 엔진을 개발했다. 기존 프레임워크를 배제하고 AVX2 및 NEON SIMD를 직접 구현하여 연산 효율을 극대화했다.

TinyLlama 모델을 1비트 포맷으로 압축하여 350MB 미만의 RAM 점유율과 150 TPS 이상의 추론 속도를 달성했다. SpQR 기법을 적용하여 양자화로 인한 지능 손실을 방지하고 100% 지능 유지를 구현했다.

메모리 대역폭 병목을 해결하기 위해 Prefill-GEMM 배치 처리와 선형 메모리 접근 방식을 도입했다. Rayon을 사용하여 모든 CPU 코어에 연산을 분산하고 정적 링 버퍼를 통해 메모리 할당 오버헤드를 제거했다.

용어 해설

1비트 양자화(1-bit Quantization): — 모델 가중치를 -1, 0, 1의 세 가지 값으로 제한하여 메모리 사용량을 극단적으로 줄이는 양자화 기법. 연산 복잡도를 낮추어 CPU 환경에서도 빠른 추론을 가능하게 한다.
단일 명령 다중 데이터(SIMD): — 단일 명령어로 여러 데이터를 동시에 처리하는 병렬 연산 기술. AVX2나 NEON 같은 명령어 세트를 활용하여 CPU의 연산 성능을 극대화한다.
일반 행렬 곱셈(GEMM): — General Matrix Multiply의 약자로, 신경망 추론의 핵심 연산. LLM의 가중치와 입력 데이터 간의 행렬 곱셈을 효율적으로 수행하는 것이 추론 속도의 관건이다.
SpQR: — 희소 행렬을 효율적으로 처리하기 위한 양자화 기법. 중요한 가중치를 보존하면서 나머지 가중치를 압축하여 모델의 지능 손실을 최소화한다.

코드 예제

text

[1] Initializing LeviathanEngine... LOADING EXACT FILE: qwen_fluent.leviathan2 ADAPTER MOUNTED: model.layers.19.self_attn.q_proj.adapter_weight ADAPTER MOUNTED: model.layers.19.self_attn.k_proj.adapter_weight ADAPTER MOUNTED: model.layers.19.self_attn.v_proj.adapter_weight ADAPTER MOUNTED: model.layers.19.self_attn.o_proj.adapter_weight ADAPTER MOUNTED: model.layers.19.mlp.gate_proj.adapter_weight ADAPTER MOUNTED: model.layers.19.mlp.up_proj.adapter_weight ADAPTER MOUNTED: model.layers.19.mlp.down_proj.adapter_weight

엔진 부팅 시 동적 INT8 어댑터가 병렬 추론 패스에 로드되는 로그

언급된 도구

Rayon추천

병렬 처리 라이브러리

TinyLlama중립

1비트 테스트 모델

Qwen중립

2비트 테스트 모델