[분석/검증-3] BitNet.cpp ARM 커널 재작성 : 스칼라 폴백구현을 통한 최종 검증

1. 개요 및 이전 상황 요약

https://uno-kim.tistory.com/458

[분석/검증-2] BitNet.cpp 텐서 연산 붕괴 현상 분석 : ARM 아키텍처 점곱 가속기의 하드웨어 결함 교

험을 통해, Exynos 프로세서에서 NEON 가속기를 강제 비활성화(-march=armv8-a+nosimd)할 경우 컴파일 에러가 발생함을 확인했었습니다.

원인은 소스 코드(ggml-bitnet-mad.cpp) 내에 NEON이나 AVX2 같은 SIMD(단일 명령 다중 데이터 처리) 가속기를 사용할 수 없을 때 동작해야 할 스칼라 폴백(Scalar Fallback) 코드가 누락되어 있었기 때문입니다.

용어 설명: 스칼라 폴백(Scalar Fallback)

 

• SIMD가 한 번에 여러 데이터를 병렬 처리하는 '가속' 연산이라면, Scalar는 데이터를 한 번에 하나씩 순차적으로 처리하는 가장 기본적인 CPU 연산 방식을 뜻합니다.
• 하드웨어 가속기를 지원하지 않는 환경을 대비해, 순수 C++ 기본 문법만으로 동일한 결과를 내도록 작성해 두는 '안전망(대체 수단)' 코드를 스칼라 폴백이라고 부릅니다.
• 그래서 이전 포스팅에서 벽돌을 하나하나 나른다라고 표현했습니다.

본 포스팅에서는 누락된 스칼라 폴백 코드를 C++로 직접 구현하고 커널에 이식하여,

"가속기를 끈 상태의 순수 연산에서는 Exynos가 정상적인 텍스트를 출력하는가?"를 최종 검증하겠습니다.

---

2. 데이터 구조 분석 및 역공학 (Reverse Engineering)

커널을 새로 작성하기 위해서는 먼저 BitNet의 1.58비트 양자화 포맷(i2_s)이 메모리 상에 어떻게 압축(Packing)되어 있는지 파악해야 합니다.

기존 AVX2 코드를 분석하여 데이터 구조를 역산했습니다.

grep -A 5 "struct block_i2_s" 3rdparty/llama.cpp/ggml/src/ggml-quants.h || grep -A 5 "struct block_i2_s" src/ggml-bitnet-mad.h

보통 llama.cpp는 양자화 데이터를 복잡한 struct로 관리하지만, 코드를 검색해 본 결과 1.58비트(i2_s)는 별도의 구조체 없이 날것의 바이트 배열을 사용하고 있음을 확인했습니다

cat -n src/ggml-bitnet-mad.cpp | sed -n '198,210p;226,234p'

누락된 C++ 폴백 코드를 작성하기 위해, 기존에 구현되어 있던 인텔 AVX2 어셈블리(Intrinsic) 코드를 뜯어 역공학(Reverse Engineering)을 진행했습니다.

• 200번 줄: 데이터를 uint8_t(1바이트) 단위로 읽어 들입니다.
• 208번 줄: 0x03 (이진수 00000011) 비트 마스크를 세팅합니다.
• 226~228번 줄: 데이터를 2비트, 4비트, 6비트씩 시프트(srli)하여 밀어냅니다.

세 가지 단서를 통해 1바이트 안에 4개의 데이터가 2비트씩 쪼개져 들어가 있다'는 패킹 규칙을 완벽하게 파악할 수 있었습니다.

2.1. 1.58비트 메모리 패킹 구조

분석 결과, 1.58비트(-1, 0, 1) 데이터는 별도의 복잡한 구조체(Struct) 없이, 가장 기본적인 1바이트 자료형인 uint8_t 배열에 압축되어 있었습니다.

 

• 1 Byte (8 bits) = 2 bits × 4개
• 즉, 1바이트 공간 안에 4개의 가중치(Weight) 데이터가 2비트씩 쪼개져 들어가 있는 형태다.
• 디코딩을 위해서는 0x03 (이진수 00000011) 비트 마스크를 사용하여 2비트 단위로 데이터를 추출해야 한다.

---

3. 스칼라 폴백(Scalar Fallback) C++ 코드 구현

따라서, SIMD 명령어를 전혀 사용하지 않고 오직 C++의 비트 연산만으로 행렬의 점곱을 수행하는 코드를 작성했습니다.

 

#else
    // ====================================================================
    // 순수 C++ 스칼라 폴백 (Scalar Fallback for ARM/Exynos)
    // ====================================================================
    const uint8_t * x_ptr = (const uint8_t *)vx;
    const int8_t  * y_ptr = (const int8_t  *)vy;

    // PC에서 변환된 GGUF는 무조건 QK_I2_S = 128 로 패킹되어 있습니다.
    const int qk = 128; 
    const int nb = n / qk; 

    for (int row = 0; row < nrc; row++) {
        int sumi = 0;
        const uint8_t * x_row = x_ptr + row * (bx / 4);

        for (int b = 0; b < nb; b++) {
            const uint8_t * px = x_row + b * 32;     // 1블록(128개 텐서) = 32 바이트
            const int8_t  * py = y_ptr + b * 128;    // 1블록 = 128 활성화 값

            for (int k = 0; k < 32; k++) {
                uint8_t xb = px[k];

                // AVX2의 _mm256_srli_epi16 추출 순서와 100% 동일하게 분할
                int v0 = (xb >> 6) & 0x03; // 비트 6,7
                int v1 = (xb >> 4) & 0x03; // 비트 4,5
                int v2 = (xb >> 2) & 0x03; // 비트 2,3
                int v3 =  xb       & 0x03; // 비트 0,1

                // [핵심] AVX2 커널은 -1을 빼지 않고 0, 1, 2를 그대로 곱합니다!
                // [핵심] AVX2는 32칸씩 건너뛰는(Interleaving) 배열 구조를 사용합니다!
                sumi += v0 * py[k];
                sumi += v1 * py[k + 32];
                sumi += v2 * py[k + 64];
                sumi += v3 * py[k + 96];
            }
        }
        // [핵심] 스케일은 상위 프레임워크(ggml_mul_mat)가 처리하므로 그대로 실수 반환!
        s[row] = (float)sumi; 
    }
#endif

4개의 핵심 함수에 스칼라 폴백 코드를 이식하곘습니다.

• ggml_vec_dot_i2_i8_s_1x1

xb >> 6 등의 비트 시프트 연산자를 통해 압축된 비트를 우측으로 밀어내고, & 0x03 마스킹을 통해 정확히 하위 2비트 값만 걸러내는 방식입니다.

이를 통해 SIMD 하드웨어 없이도 CPU의 기본 논리 연산기만으로 데이터를 정확하게 해석할 수 있습니다.

https://github.com/uno-km/BitNet/tree/main

GitHub - uno-km/BitNet: BitNet inference framework with ARM/Exynos Scalar Fallback implementation.

코드는 제가 포크해서 직접 수정을 해서 저장헀습니다.

QK_I2_S 정의 수정, ggml_vec_dot_i2_i8_s 수정을 추가적으로 더 했습니다.

---

4. 커널 이식 및 빌드

작성한 코드를 src/ggml-bitnet-mad.cpp 파일 내부에 누락되어 있던 점곱 함수 4곳(ggml_vec_dot_i2_i8_s_1x1 등)의 #endif 직전(AVX2와 NEON 분기문 끝부분)에 #else 구문으로 이식했습니다.

이후, 설정했던 가속기 완전 차단 플래그를 유지한 채 수동 빌드를 재실행했습니다.

++빌드중 오류발생건

우리가 직접 수정한 ggml-bitnet-mad.cpp는 [ 8%] 단계에서 에러 없이 성공적으로 컴파일되었습니다!

그러나 llama.cpp 개발진이 "ARM(aarch64) 프로세서라면 당연히 NEON 가속기가 있을 것"이라고 가정하고 코드를 짰기 때문에

우리가 강제로 NEON을 꺼버리자, 가속기를 전제로 정의되던 변수(ROW_BLOCK_SIZE 등)들이 증발해 버렸고, 이를 참조하던 다른코드들에서 오류가 났습니다.

# 에러가 발생한 파일의 맨 윗줄에 누락된 블록 사이즈 정의를 추가합니다.
sed -i '1i #define ROW_BLOCK_SIZE 1\n#define COL_BLOCK_SIZE 1' 3rdparty/llama.cpp/ggml/src/ggml-aarch64.c

무사히 빌드중...너무다행이다 ㅠㅠ

모든 명령어

cd ~/workspace
rm -rf BitNet

# 내 깃허브에서 클론 & 서브모듈(llama.cpp) 채워넣기
git clone https://github.com/uno-km/BitNet.git
cd BitNet
git submodule update --init --recursive

# 필수 패키지 설치
pip install -r requirements.txt
pip install torch transformers sentencepiece
pip install -e 3rdparty/llama.cpp/gguf-py

# 모델 다운로드 및 변환 (시간이 조금 걸립니다)
python setup_env.py --hf-repo microsoft/BitNet-b1.58-2B-4T --quant-type i2_s

# 헤더 파일 생성 (파이썬이 안 해줬을 경우를 대비)
python utils/codegen_tl1.py --model bitnet_b1_58-3B --BM 160,320,320 --BK 64,128,64 --bm 32,64,32

# ggml-aarch64.c 파일 맨 윗줄에 더미 값(1) 강제 주입
sed -i '1i #define ROW_BLOCK_SIZE 1\n#define COL_BLOCK_SIZE 1' 3rdparty/llama.cpp/ggml/src/ggml-aarch64.c

# build 폴더 삭제
rm -rf build

# 완벽한 스칼라(NEON OFF) 세팅 주입
cmake -B build \
  -DGGML_NEON=OFF \
  -DLLAMA_NEON=OFF \
  -DGGML_NATIVE=OFF \
  -DLLAMA_NATIVE=OFF \
  -DGGML_LLAMAFILE=OFF \
  -DLLAMA_LLAMAFILE=OFF \
  -DCMAKE_C_COMPILER=clang \
  -DCMAKE_CXX_COMPILER=clang++ \
  -DCMAKE_C_FLAGS="-march=armv8-a+nosimd -U__ARM_NEON" \
  -DCMAKE_CXX_FLAGS="-march=armv8-a+nosimd -U__ARM_NEON"

# 1코어로 안전하게 빌드
cmake --build build -j 1

python run_inference.py \
  -m models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf \
  -p "The capital of France is" \
  -n 50 \
  -temp 0.0 \
  -t 1

어?!

드디어 제대로 출력되었습니다!!!!!!!!!

그토록 듣고싶었던 파리...

이제 이 긴 여행도 끝에 다달았습니다.ㅠㅠㅠㅠ

다음 포스팅은 결과보고 작성으로 뵙겠습니다!!!!!!