Low Bit ARM kernels

데스크탑 / 폰 / 임베디드 — NVIDIA GPU 가 없는 환경. 이 자리에서 LLM 추론을 돌리려면 거의 모든 결정이 다르게 내려진다. CPU 위주 compute, 적은 메모리, 강력한 cache hierarchy, 저전력 제약. Scott 의 일은 그 자리에 맞는 low-bit kernel set 을 짜는 것.

강의의 첫 측정 — Llama 3.1 8B 가 M1 Pro 맥북에서 토큰 당 약 50ms. 4-bit weight + INT8 activation. 이게 llama.cpp 의 길을 따라가지 않고 PyTorch 생태계 안에서 (torchao + ExecuTorch) 직접 짠 결과. macOS / Linux / iOS 의 같은 stack 으로 동작.

강의 끝에 손에 잡혀야 할 자산 — (1) ARM CPU 의 compute 단위(NEON / SVE) 와 GPU SM 의 직관적 매핑, (2) low-bit kernel 의 “register-level tiling + cache-level tiling” 이라는 두 단계 구조, (3) ExecuTorch 의 custom op 가 어떻게 모델 안에 들어가는지, (4) 같은 1~8 bit 범위에서 정확도와 latency 의 trade-off 가 GPU 와 어떻게 다른지.

데스크탑/모바일 환경에서 LLM 의 자연스러운 compute target 이 CPU 인 이유는 단순하다 — 가장 통합된, 가장 큰, 가장 호환되는 unit 이기 때문.

강의에서 Scott 이 명시한 한 줄 — “CPU 가 가장 portable 하고 가장 통합된 target. 모바일 GPU/NPU 는 좋지만 fragmentation 이 심해서 한 stack 으로 모든 device 를 cover 하기 어렵다.” 그래서 ExecuTorch 는 CPU backend 를 main 으로 두고, GPU/NPU 는 가능한 device 에서 추가로 넘긴다.

NVIDIA GPU 는 SIMT (한 thread 가 한 element, 32 thread = warp). ARM 은 SIMD (한 thread 가 여러 element 동시). 이 차이가 mental model 의 첫 번째 갈림길.

• NEON (ARMv8) — 128-bit fixed-width vector. fp32 lane 4 개, int8 lane 16 개, int16 lane 8 개. 모든 ARM 기기에 있다.
• SVE / SVE2 (ARMv9) — vector 폭이 가변 (128~2048 bit). compile 한 번으로 다양한 폭의 hardware 에서 동작. ARMv9 에서 점차 표준.
• I8MM 확장 — int8 matrix multiply (smmla). 한 instruction 이 4×4 int8 GEMM block 을 한 번에. low-bit kernel 의 중심.

강의에서 Scott 이 짚은 사실 — FMLA (vector FMA) 는 fp16/fp32 lane 별 곱셈+누적이고, SMMLA 는 int8 matrix 곱. low-bit kernel 의 inner loop 는 거의 SMMLA. dequant 결과를 int8 형태로 쥐고 있다가 SMMLA 로 던진다.

decode 는 batch=1 → activation 이 vector 한 줄, weight 가 matrix. GEMV (matrix-vector). prefill 만 GEMM. 모바일 LLM 의 latency 는 거의 GEMV 가 결정.

강의에서 Scott 이 명시한 — “mcal=1, ncal=8, kcal=16 가 GEMV 의 표준 tile.” M=1 인 이유는 decode batch=1, N=8 은 NEON int32 acc 4 개 × 2 SMMLA, K=16 은 SMMLA 한 발행이 받는 단위.

CPU kernel 의 tiling 은 GPU 와 다르게 명확한 hierarchy 가 있다. register tile (가장 안쪽), L1 cache tile (중간), thread tile (가장 바깥). Scott 의 ARM kernel 도 같은 구조.

이 hierarchy 가 GPU 의 “grid → block → warp → thread” 와 닮아 보이지만 의미가 다르다 — CPU 에서는 cache layer 의 명시성 이 강하다. L1 / L2 / L3 의 크기가 분명하고, working set 이 cache 안에 들어가는지 안 들어가는지가 latency 를 직접 결정.

weight 는 보통 한 번 load 되면 그 thread 동안 여러 번 reuse. activation 은 한 번 load + 한 번 use. 그래서 weight 의 cache 친화성 이 더 결정적. weight 를 미리 packing 단계에서 cache-friendly layout 으로 적어둔다.

L034 의 GPU 와 마찬가지로, ARM low-bit 도 핵심은 group-wise scale. 다만 group_size 의 단위가 더 작다 (32~128). 그리고 ARM 쪽은 zero point (asymmetric quant) 도 더 자주 쓴다.

• symmetric — scale 만. zero point = 0. 빠르지만 weight 분포가 한쪽으로 치우치면 손해.
• asymmetric — scale + zero point. 한 group 마다 fp16 scale + int8 zp. dequant 코드가 한 줄 늘어나지만 정확도가 안정.
• group_size 표준 — 32, 64, 128. 32 가 정확도 best, latency cost 약 5~10%. ExecuTorch 표준은 보통 32.

강의에서 흥미로운 사실 — 3-bit 가 4-bit 와 거의 같은 정확도를 주는 케이스가 의외로 많다. weight 의 effective bits 를 1 bit 줄이는 게 모바일 메모리 절약에 큰 의미. 다만 unpacking 코드가 4-bit 보다 복잡해서 latency 는 비슷하거나 약간 더 느림.

kernel 이 빠른 것만으로는 의미가 없다 — 실제 모델 안에 들어가야. 강의 후반부는 ExecuTorch 의 custom op pipeline 위에 어떻게 이 kernel 들이 박히는지를 깐다.

핵심은 두 가지 — (a) kernel 이 일반 C++ 함수, 즉 추가 컴파일러나 DSL 안 필요. torch.ops.{namespace} 로 등록되며 PyTorch / ExecuTorch 모두에서 호출 가능. (b) 병렬화는 pthreadpool (XNNPACK 의 thread pool wrapper) 를 그대로 빌려 씀. ExecuTorch / PyTorch / XNNPACK 가 같은 thread pool 을 공유.

// torchao 의 ARM kernel — C++ 그대로
struct UKernelConfig {
  int mcal, ncal, kcal;             // register tile
  int n_tiles_per_thread;
  KernelFn kernel;                  // 함수 포인터
  UnpackFn unpack;                  // 1~8 bit 별
};

// 24 개 GEMV kernel + 8 개 unpacking routine
// 의 lookup table — bitwidth × tile
extern UKernelConfig kernels[8][3];

강의에서 보여준 측정 — 같은 모델 (Llama 3.1 8B), 같은 quant (4-bit weight + dynamic INT8 act, group=32), 같은 stack (torchao + ExecuTorch + ARM kernels) 의 device 별 latency.

두 가지 사실. 첫째 — 맥북에서도 사용 가능한 속도가 나온다. 20 tok/s 면 사람이 읽기 따라잡는 속도. 둘째 — iPhone 은 메모리가 hard limit. 8GB RAM 의 디바이스에서 5GB 모델은 OS 와의 경쟁에서 swapping 위험.

3B 클래스 모델 (Llama 3.2 3B) 으로 가면 iPhone 에서도 20+ tok/s 가 나오는 영역. 모바일 LLM 의 sweet spot 이 1B~4B 인 이유 — 메모리 limit + latency 의 두 제약이 합쳐서.

표가 보여주는 것 — 두 hardware 는 완전히 다른 trade-off space 위에 있다. cloud GPU 는 throughput 게임(batching, prefix sharing, 대량 동시 요청), edge 는 latency / privacy / offline 게임(batch=1, 모델 작아야 함). 같은 quantization 이라도 hardware 의 다른 제약이 코드를 다르게 펴게 만든다.

손에 새기기 — 실습 시퀀스

1. NEON intrinsic 한 줄 — vfmaq_f32, vmmlaq_s32. C++ 한 함수로 vectorized fp32 dot product. 결과를 일반 코드와 비교.
2. SMMLA loop — int8 weight + int8 act 의 8×16 tile 한 발행. 결과를 reference 와 비교. 이게 ARM low-bit 의 “tensor core hello world”.
3. register tile 만 단독 — pthreadpool 없이 single thread 에서 GEMV inner kernel. M1 Pro 에서 ~10 tok/s 는 single thread 에서 이미 가능.
4. parallel_1d 추가 — pthreadpool 띄우고 N 차원 분할. core 수만큼 throughput 곱 — 8 core 면 ~80 tok/s 까지 toy GEMV 에서.
5. torchao 의 quantize → packing — Llama 3 한 layer 를 quantize 해서 ARM kernel 에 넘긴다. round-trip 으로 정확도 검증.
6. ExecuTorch export — torch.export + ExecuTorch binary. 그 binary 를 iPhone 또는 raspberry pi 로 옮겨 동작 확인.
7. llama.cpp 와 같은 모델 비교 — 같은 weight, 같은 quant 에 대해 두 stack 의 latency. ARM kernel 이 어디서 더 빠른지/느린지.

• SVE2 의 본격 활용 — ARMv9 가 모바일 SoC 에 본격 보급되면 가변 vector 폭의 이득은 어디서 나오는가. NEON 만 쓰는 코드와 latency 차이.
• NPU / DSP 통합 — Apple Neural Engine, Qualcomm Hexagon. ExecuTorch 가 이 backend 들을 표준화하는 방향. 정확히 어느 op 가 NPU 로 갈 수 있는가.
• P-core / E-core 비대칭 work scheduling — 강의 시점의 parallel_1d 는 균등 분할. 그게 정확히 얼마나 손해를 보고 있는지.
• 모바일 GPU (Mali/Adreno/Apple) 와의 fallback — 같은 모델을 device 에 따라 CPU/GPU/NPU 로 분기. 그 dispatch 결정의 휴리스틱은 어디에 있는가.
• bit-width 와 정확도의 거의 정확한 곡선 — 1, 2, 3, 4, 5, 6, 7, 8 bit 별 perplexity 곡선. 어떤 모델에서 3-bit 가 4-bit 와 거의 같은지.