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

지난 편에서 비동기 softmax를 커스텀 연산자 OptSoftmax로 구현하고 긴 디버깅 끝에 정확도를 맞췄습니다. 그런데 정확도를 확보한 뒤 성능 프로파일을 보자, 커스텀 연산자 접근 자체에 더 근본적인 한계가 있음을 알게 됐습니다. 이번 편은 그 한계를 확인하고 최적화를 넣는 위치를 바꾼 이야기입니다.

1. 구조적 한계: 커스텀 연산자는 HMX 할당을 깨뜨린다

QNN 컴파일러는 MatMul을 HMX(행렬 유닛)에 할당하도록 설계돼 있는데, 커스텀 연산자를 만들면 해당 연산이 컴파일러의 최적화 경로에 등록되지 않아 HMX 활용도가 떨어졌습니다.

두 가지 융합 범위를 실험했습니다. 처음에는 attention 전체를 커스텀 연산자로 묶었는데, 그러자 내부의 MatMul까지 커스텀 영역으로 흡수되어 컴파일러가 이를 HMX로 할당하지 못하고 HVX로 내려보냈습니다. 결론은 MatMul을 건드리면 안 된다는 것이었습니다. 다음으로 MatMul은 그래프에 남기고 Softmax와 마스크 덧셈만 커스텀화하자 활용도가 회복됐지만, 여전히 원본 모델 수준에는 미치지 못했습니다.

두 실험이 가리키는 결론은 분명했습니다. 커스텀 연산자를 등록하는 한, 원본 모델의 HMX 활용도를 따라갈 수 없다. 컴파일러가 인식하는 표준 연산자 경로를 벗어나는 순간 하드웨어 가속 효율이 떨어지고 불필요한 메모리 복사가 생기기 때문입니다.

커스텀 연산자의 융합 경계는 그래프의 미학이 아니라 유닛 할당(HMX vs HVX)이 결정합니다. 그리고 그 경계를 아무리 잘 그어도, 커스텀 영역이 존재하는 한 컴파일러의 Graph Matching·Fusion이 닿지 못하는 사각지대가 남습니다.

2. 방향 전환: Export 단계에서 async softmax 적용

그래서 접근을 바꿨습니다. 커스텀 연산자를 런타임에 등록하는 대신, ONNX export 이전 단계—PyTorch modeling.py의 attention 구현—에 비동기 softmax를 직접 넣고 그대로 export하는 방식입니다. 이렇게 하면 최종 그래프는 여전히 표준 연산자들로 구성되므로, 컴파일러의 fusion과 HMX 할당이 온전히 유지됩니다.

한 가지 오해하기 쉬운 부분이 있습니다. 여기서의 KV 캐시 분할은 export가 끝난 ONNX 그래프를 나중에 잘라내는 작업이 아닙니다. async softmax의 chunk 분할은 attention이 KV 캐시를 읽는 바로 그 지점에서 일어나야 하므로, PyTorch modeling.py의 attention 안에서 KV 캐시를 직접 split하고 chunk별 부분합·정규화에 필요한 텐서를 추가한 다음 export해야 합니다. 그래야 export된 그래프에 이 구조가 표준 연산자로 온전히 펼쳐지고, 컴파일러가 곧바로 fusion·HMX 할당을 적용할 수 있습니다.

접근 1 · Custom Op 등록접근 2 · Export 단계 적용 (채택)
적용 위치QNN 그래프에 OptSoftmax 커스텀 연산자 등록PyTorch modeling.py attention에 async softmax 삽입 후 ONNX export
HMX 활용도컴파일러가 커스텀 영역을 인식 못해 MatMul이 HMX에서 이탈 · 원본 미달표준 연산자 유지 → fusion 보존 · 91 → 94
비용 / 과제커널 구현·검증, Genie 등록(SDK 내부 이해) 등 숨은 부담커스텀 등록·검증 부담 제거 · TPS 약 10 하락은 과제
3. 결과: HMX 활용도와 TPS의 어긋남

결과는 설계 의도대로였습니다. 표준 경로를 지킨 덕분에 HMX 활용도가 원본 91에서 94로 올랐습니다. 커스텀 연산자로는 원본을 따라가지도 못했던 지표를, 최적화를 넣는 위치만 바꿔 오히려 넘어선 것입니다. 다만 활용도 상승이 곧 처리량 향상을 뜻하지는 않았습니다. 같은 구성에서 TPS는 약 10만큼 낮아졌습니다.

실측 프로파일이 이를 뒷받침합니다.

지표원본 모델Export 단계 적용
HMX 이용률91.2 %94.9 %
HTP 노드 수8,60616,066
디코딩 TPS기준약 −10

특기할 점은, 단일 그래프 프로파일러상의 지표(graph execute 5,652 → 4,836 µs, 193.6 inf/s)는 오히려 개선됐는데도 실제 디코딩 TPS는 낮아졌다는 것입니다. KV 캐시 분할을 표준 연산자로 펼치면서 HTP 노드 수가 8,606 → 16,066으로 거의 두 배가 됐고, 이때 늘어난 split·concat·reduction이 토큰마다 반복되는 디코드 루프에서 오버헤드로 나타난 것으로 의심됩니다.

이 어긋남 자체가 이번 작업의 중요한 인사이트입니다. HMX 활용도는 "유닛이 얼마나 바쁜가"를 말할 뿐, "토큰이 얼마나 빨리 나오는가"를 보장하지 않습니다. 활용도를 회복시켰다는 사실보다, 활용도와 TPS를 함께 끌어올리는 지점을 찾는 일이 남은 과제라는 점이 더 중요합니다. 현재는 이 TPS 격차의 원인을 분해하고 있으며, 이를 좁히는 것을 다음 단계로 두고 있습니다.

4. Genie 통합의 실체

커스텀 연산자 접근이 값비쌌던 또 다른 이유는 추론기 통합에 있었습니다. Genie는 공식적으로 커스텀 연산자 추가를 지원하지 않으므로, op package를 그래프 정의가 아니라 백엔드 초기화 코드에서 직접 등록해야 합니다. 백엔드 핸들이 준비되는 시점에 backendRegisterOpPackage를 호출해 op package(.so)를 주입하도록 QnnApi.cpp를 수정하는 식인데, 이는 Genie SDK의 내부 구조를 상당히 이해해야 가능한 작업이었습니다.

이 통합 비용은 export 단계 접근이 매력적인 또 하나의 이유였습니다. 최적화를 표준 연산자로 표현해 두면 Genie의 디코드 루프를 그대로 재사용할 수 있고, 커스텀 등록이라는 유지보수 부담 자체가 사라지기 때문입니다.

5. 결론: 최적화는 "무엇을"만큼 "어디에"가 중요하다

이번 이식의 가장 큰 교훈은 알고리즘의 내용이 아니라 그것을 넣는 위치에 있었습니다. 같은 비동기 softmax라도 커스텀 연산자로 등록하면 컴파일러의 HMX 할당을 깨뜨리고, export 단계에서 표준 연산자로 표현하면 오히려 활용도를 끌어올립니다.

정리하면 세 가지입니다.

  • 컴파일러가 지원하는 Graph Matching·Fusion을 살리려면 가급적 표준 연산자를 유지하고, 그래프 수정 전에 SDK가 지원하는 연산 조합으로 대체 가능한지부터 검토해야 합니다.
  • 커스텀 연산자는 커널 구현·검증(양자화, HVX 순서, KV 캐시 누적)과 Genie 등록 모두에서 예상 밖의 비용을 요구합니다.
  • HMX 활용도와 TPS는 같은 방향으로 움직이지 않으므로, 활용도 회복을 성과의 종착점이 아니라 처리량 개선을 위한 출발점으로 삼아야 합니다. 지금은 HMX 활용도 이득을 유지하면서 TPS 격차를 좁히는 방향으로 작업을 이어가고 있습니다.

긴 글 읽어주셔서 감사합니다!!

행렬 곱에 특화된 Hexagon의 연산 유닛입니다. 큰 GEMM에서 높은 처리량을 내며, 커스텀 연산자가 MatMul을 흡수하지 않도록 그래프를 설계해야 이 유닛이 제대로 활용됩니다.

벡터 연산을 담당하는 유닛으로, softmax의 exp·sum 같은 원소별 연산에 쓰입니다. 데이터 순서가 하드웨어 특성에 좌우됩니다.

컴파일러가 표준 연산자들의 특정 패턴을 인식해 하나의 가속 커널로 묶는 최적화입니다. 커스텀 연산자로 경로를 벗어나면 이 최적화가 적용되지 않습니다.

참고 문헌
  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.