[Engineering] FlashDecoding++ 이식기 (2) — 커스텀 연산자와 디바이스 디버깅

안녕하세요, 옵트에이아이 박성재입니다.

지난 편에서 정리한 비동기 softmax를 Qualcomm Hexagon NPU에 올리기 위해, 가장 직관적인 출발점으로 attention을 커스텀 연산자(Custom Op)로 직접 구현했습니다. 이번 편은 그 구현과, 무엇보다 그 뒤에 이어진 디바이스 디버깅의 기록입니다. 결론부터 말하면 커스텀 커널은 구현보다 검증이 훨씬 무거웠습니다.

1. 첫 시도: attention 전체를 하나의 커스텀 연산자로

가장 먼저 손댄 것은 attention 전체 — MatMul → Add(Mask) → Softmax → MatMul — 를 통째로 하나의 커스텀 연산자로 묶는 것이었습니다. 중간 텐서를 밖으로 내보내지 않으면 메모리 왕복이 줄어들 것이라는, 단순하고 자연스러운 기대였습니다.

기대대로 DRAM read는 크게 줄었습니다. 그런데 지연 시간은 오히려 나빠졌습니다.

원인은 HMX 할당이 깨진 데 있었습니다. QNN 컴파일러 입장에서 커스텀 연산자 내부는 들여다볼 수 없는 블랙박스입니다. attention의 두 MatMul까지 이 블랙박스 안에 넣어버리자, 컴파일러는 이들을 행렬 가속기인 HMX에 매핑하지 못하고 전부 HVX(벡터 유닛)로 떨어뜨렸습니다. HMX 이용률은 91% → 31%로 급락했고, HVX로 쏟아진 연산 탓에 HTP 노드 수는 8,606 → 30,961개로 폭증했습니다. 중간 텐서를 덜 쓴 덕에 DRAM read는 295M → 164M로 줄었지만, MatMul을 HMX에서 잃은 손해가 이를 압도해 처리량이 159.9 → 113.8 inf/s로 오히려 떨어졌습니다.

교훈은 분명했습니다. MatMul은 HMX에 남겨두어야 한다. 그래서 융합 범위를 좁혔습니다 — 두 MatMul은 네이티브 QNN 연산으로 두어 HMX가 처리하게 하고, 그 사이의 메모리 바운드 구간인 마스크 덧셈과 softmax만 하나로 묶기로 했습니다. 이 판단의 결과가 다음 절의 OptSoftmax입니다.

2. OptSoftmax 커스텀 연산자

기존 attention의 기본 구조는 다음과 같습니다.

scores = Q × Kᵀ
scores += mask
prob    = softmax(scores)
output  = prob × V

이 구조는 softmax를 위해 [B, H, Q, K] 전체를 메모리에 유지해야 하며, KV 길이가 커질수록 메모리 대역폭과 VTCM 사용량이 늘고 global dependency 때문에 스트리밍이 불가능합니다. 그래서 softmax를 chunk 단위로 분할하고 global normalization으로 재구성한 뒤(1편의 비동기 아이디어를 2-pass로 옮긴 것), 마스크 덧셈과 softmax를 하나의 커스텀 연산자 OptSoftmax로 융합했습니다.

변경 전:  MatMul → Add(Mask) → Softmax → MatMul
변경 후:  MatMul → OptSoftmax(Custom Op) → MatMul

*변경 전 — MatMul·Add(Mask)·Softmax·MatMul이 개별 노드로 나열된 표준 그래프*

*변경 후 — Add와 Softmax가 OptSoftmax 하나로 융합된 그래프*

커널 내부에서는 각 블록이 online softmax처럼 최댓값을 매번 탐색하는 대신 사전 계산된 통합 최댓값 fdpp_phi를 사용하고, VTCM에 버퍼 두 개(row_buf[2])를 두어 현재 루프의 연산과 다음 루프의 로딩을 겹쳤습니다.

*두 버퍼를 번갈아 쓰며 로딩과 연산을 상시 중첩시킨다*

3. 반복되는 토큰: 디버깅의 시작

이식 초기 가장 끈질겼던 증상은 같은 단어가 무한히 반복되는 출력이었습니다. 원인을 좁히기 위해 판단 기준을 세웠습니다. 첫 토큰부터 틀리면 encodings·tokenizer를, 두 번째 토큰부터 무너지면 KV 캐시·mask·RoPE·attention을 의심하는 식입니다. 우리 케이스는 첫 토큰부터 어긋났으므로 encodings와 양자화를 먼저 파고들었습니다.

4. CPU 백엔드로 로직부터 검증

먼저 host 환경에서 CPU 백엔드(libQnnCpu.so)로 FP32 로직을 검증했습니다. qairt-converter에 --input_layout 옵션이 빠져 LLM의 4D 텐서가 비전용 이미지 포맷으로 강제 변환되며 토큰 순서가 섞이던 문제를 NONTRIVIAL 지정으로 바로잡자, 데이터 개수·마스크 패딩·argmax가 모두 정상으로 나왔습니다. CPU 코드 자체는 문제가 없었고, 원인은 HTP의 양자화 경로였습니다.

5. 양자화가 마스크의 의미를 바꾸고 있었다

로그를 분석하자 마스크 값이 의도한 값이 아니라 -3000 근처로 뭉개져 masking이 제대로 작동하지 않았습니다. QNN HTP는 텐서를 uint16 정수로 저장하지만 softmax의 exp·sum·divide는 float 영역에서 계산해야 정확합니다. 스코어와 마스크의 scale·offset을 올바르게 반영하도록 dequant/quant 경로를 고치자 마스킹이 정상화됐습니다.

uint16 입력 → dequantize → float → softmax → quantize → uint16 출력
real = (quant - zero_point) × scale

6. HVX의 물리적 동작이 인덱스를 뒤섞고 있었다

마스크를 고친 뒤에도 확률 값은 거의 동일한데 상위 확률의 인덱스만 몇 칸씩 밀리는 현상이 남았습니다. 원인은 uint16을 uint32로 zero-extend하는 intrinsic(Q6_Wuw_vzxt_Vuh)의 even/odd 분리 동작이었습니다. 128바이트 벡터에 담긴 uint16 64개를 uint32로 확장하면 256바이트가 되어 벡터 두 개로 나뉘는데, HVX는 셔플 네트워크 구조상 이를 짝수·홀수 인덱스로 갈라 배치합니다. 하드웨어 배선상 가장 단순한 동작이지만, softmax 출력의 순서를 그대로 흔들었습니다.

*zero-extend 시 uint16이 짝수·홀수 레인으로 갈라지며 출력 순서가 어긋난다*

7. 작은 편향이 KV 캐시로 누적되다

순서를 되돌려 argmax를 맞춘 뒤에도 추론이 완전히 안정되지 않아, 원본과 커스텀 모델의 프리필 KV 캐시를 층위별로 비교했습니다. Layer 0은 정상이었지만 중간 레이어부터 value가 어긋났고, iteration이 진행될수록 악화됐습니다. 초기 hidden의 코사인 유사도는 반복이 거듭되며 0.990 → 0.965 → 0.561로 떨어졌습니다. OptSoftmax의 미세한 편향이 오염된 KV 캐시로 되먹임되며 눈덩이처럼 커진 것입니다.

*초기 레이어 hidden state가 iter_2에서 0.561까지 급락한다 (프리필 층위별 실측)*

커스텀 연산자에서 "값은 맞는데 결과가 틀린" 상황은 대부분 순서·양자화·유닛 할당 중 하나입니다. 스칼라 우회 경로를 하나 마련해 두면 HVX 특유의 재배열을 제거했을 때 그 변경이 실제 .so에 반영됐는지까지 함께 검증할 수 있습니다. 그러나 이 모든 검증에 든 시간 자체가, 커스텀 커널 접근의 숨은 비용이었습니다.

참고 문헌

1. Hong, K., et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
2. Qualcomm. (2026). Qualcomm AI Engine Direct (QNN) SDK Documentation.