Published on

ONNX Runtime로 LLM INT4 양자화와 지연 개선

Authors

서빙 환경에서 LLM 지연을 줄이려면 보통 두 가지 축을 동시에 만져야 합니다. 첫째는 모델 자체를 가볍게 만드는 것(양자화, KV 캐시 최적화), 둘째는 런타임 실행 경로를 빠르게 만드는 것(그래프 최적화, 커널 선택, 스레딩, I O 패턴).

이 글은 ONNX Runtime(이하 ORT) 기준으로 LLM INT4 양자화를 적용하고, 프리필(prefill)과 디코드(decode) 지연을 개선하는 과정을 실무 관점에서 정리합니다. 특히 “INT4로 줄였는데도 생각만큼 빨라지지 않는다” 같은 상황에서 어디를 의심해야 하는지 체크리스트까지 포함합니다.

왜 INT4인가: 비용 구조가 바뀐다

LLM 추론 비용은 대체로 아래로 분해됩니다.

  • 가중치 로딩 및 메모리 대역폭: 가중치가 클수록, 그리고 토큰 생성 중 반복적으로 접근할수록 병목이 됩니다.
  • 프리필 비용: 입력 시퀀스 길이에 비례해 Attention과 MLP가 한 번에 크게 실행됩니다.
  • 디코드 비용: 토큰을 한 개씩 생성하며 작은 배치로 반복 실행됩니다. 커널 런치 오버헤드, 캐시 효율, KV 캐시 접근 패턴이 중요합니다.

INT4는 가중치 메모리를 FP16 대비 대략 4분의 1 수준으로 줄여 메모리 대역폭 병목을 완화하고, CPU 환경에서는 특히 GEMM 커널이 INT4에 최적화되어 있으면 실질 속도로도 이어질 수 있습니다. GPU에서도 효과가 있지만, GPU는 커널/텐서코어 경로, 스케일링/디퀀타이즈 비용, 그래프 최적화 여부에 따라 체감이 달라집니다.

ORT에서의 INT4 양자화: 핵심 선택지

LLM INT4 양자화는 크게 다음 갈래로 나뉩니다.

  • 가중치만 INT4(Weight-only quantization)
    • 활성값(activation)은 FP16/FP32로 유지
    • 정확도 손실이 상대적으로 작고 적용이 쉬움
    • 디코드 구간에서 특히 이득을 보는 경우가 많음
  • 가중치+활성값까지 저정밀(정수)로
    • 더 공격적인 최적화 가능
    • 캘리브레이션, 정확도 검증, 커널 지원 조건이 까다로움

ORT 생태계에서 LLM은 보통 “가중치 INT4 + 활성 FP16” 조합이 현실적인 출발점입니다.

전체 파이프라인 개요

실무에서 추천하는 순서는 아래입니다.

  1. 모델을 ONNX로 내보내기(가능하면 이미 ONNX 형태의 LLM 사용)
  2. 양자화 적용(INT4 weight-only 중심)
  3. ORT 세션 옵션 튜닝(그래프 최적화, 스레딩, EP 설정)
  4. 프리필/디코드 분리 측정(단일 숫자 지연이 아니라 구간별)
  5. 정확도 및 회귀 테스트(응답 품질, 특정 프롬프트에서 붕괴 여부)

이 과정은 Stable Diffusion에서 VRAM을 줄이기 위해 xFormers나 타일링을 조합하듯, “한 방”보다는 여러 레버의 조합으로 성능을 끌어올리는 방식에 가깝습니다. 메모리 압박을 줄이는 사고방식은 아래 글과도 결이 비슷합니다.

INT4 양자화 실전: Python 예제

아래 예제는 “ONNX 모델을 로드해서 양자화한 뒤 저장”하는 기본 흐름입니다. 실제로는 모델 구조(예: Attention fused 여부, MatMul 패턴)에 따라 적용 가능한 옵션이 달라질 수 있으니, 일단 작은 샘플로 검증하고 본 모델로 확장하는 것을 권장합니다.

1) 의존성 설치

pip install onnx onnxruntime onnxruntime-tools

환경에 따라 onnxruntime-gpu 를 설치해 CUDA Execution Provider(EP)를 쓰기도 합니다.

pip install onnxruntime-gpu

2) 동적 양자화(참고)와 INT4 방향

전통적인 ORT 양자화는 INT8 동적 양자화가 널리 알려져 있습니다. INT4는 “모델 패턴과 커널 지원”이 더 중요하므로, 다음을 먼저 확인하세요.

  • 목표가 CPU인가 GPU인가
  • ORT 빌드가 INT4 커널을 충분히 활용할 수 있는가
  • 모델이 ORT가 최적화하는 MatMul/Gemm 패턴으로 정리되어 있는가

아래 코드는 구조를 보여주는 예시입니다. 실제 INT4 weight-only는 ORT 버전과 도구에 따라 API가 다를 수 있어, 본인의 ORT 버전에 맞는 툴(예: ORT용 quantization toolchain)을 확인해야 합니다.

import onnx
from pathlib import Path

# 예: 원본 ONNX
model_path = Path("./llm_fp16.onnx")
model = onnx.load(str(model_path))

# 여기서 INT4 weight-only 양자화 로직을 적용한다고 가정
# (실제 적용 함수/옵션은 ORT 버전 및 도구에 따라 달라질 수 있음)
# 예시: 특정 MatMul/Gemm 가중치 텐서를 INT4로 패킹하고 스케일을 추가

onnx.save(model, "./llm_int4.onnx")
print("saved: llm_int4.onnx")

포인트는 “코드 한 줄로 끝”이 아니라, 모델 그래프가 양자화 친화적으로 정리되어야 한다는 것입니다. 예를 들어 Linear 레이어가 MatMul + Add로 표현되어 있고, 상수 가중치가 initializer로 들어가 있어야 weight-only 양자화가 적용됩니다.

ORT 세션 튜닝: 지연을 깎는 옵션들

양자화가 끝났다면, 이제 ORT 실행 옵션으로 “실제 지연”을 줄일 차례입니다. LLM은 토큰을 반복 생성하므로, 작은 오버헤드도 누적됩니다.

1) 그래프 최적화 레벨

import onnxruntime as ort

so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
# 필요 시 프로파일링 활성화
so.enable_profiling = True

sess = ort.InferenceSession(
    "./llm_int4.onnx",
    sess_options=so,
    providers=["CPUExecutionProvider"],
)
  • ORT_ENABLE_ALL 로 기본적인 fuse, 상수 폴딩 등을 최대화합니다.
  • 프로파일링 JSON을 열어보면 어떤 노드가 시간을 먹는지 확인할 수 있습니다.

2) 스레딩과 병렬성(특히 CPU)

CPU에서 LLM은 GEMM이 지배적이라 스레딩 설정이 민감합니다.

so = ort.SessionOptions()
so.intra_op_num_threads = 8
so.inter_op_num_threads = 1
so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
  • intra_op_num_threads 는 단일 연산 내부 병렬성
  • inter_op_num_threads 는 연산 간 병렬성
  • LLM 디코드는 작은 그래프를 반복하므로, ORT_SEQUENTIAL 이 유리한 경우가 많습니다(불필요한 스케줄링 오버헤드 감소).

3) I O 바인딩으로 복사 비용 줄이기

GPU EP를 쓸 때는 특히 입력/출력 텐서 복사가 지연에 영향을 줍니다. ORT의 I O binding을 사용하면 불필요한 복사를 줄일 수 있습니다.

import numpy as np
import onnxruntime as ort

sess = ort.InferenceSession(
    "./llm_int4.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)

io = sess.io_binding()

# 예: input_ids
input_ids = np.array([[1, 2, 3, 4]], dtype=np.int64)
io.bind_cpu_input("input_ids", input_ids)

# 출력은 디바이스에 바인딩(실제 출력 이름은 모델에 맞게)
io.bind_output("logits")

sess.run_with_iobinding(io)
outputs = io.copy_outputs_to_cpu()

디코드 루프에서 토큰마다 sess.run 을 호출하는 구조라면, I O 바인딩과 메모리 재사용 여부가 체감 성능을 좌우할 수 있습니다.

프리필과 디코드: 측정 방법부터 바꿔야 한다

“평균 지연”만 보면 개선이 가려집니다. LLM은 보통 다음으로 나눠 측정합니다.

  • TTFT(Time To First Token): 프리필 비용이 크게 반영
  • TPOT(Time Per Output Token): 디코드 루프 성능

INT4는 가중치 대역폭을 줄여 TPOT에 이득이 나는 경우가 많지만, 프리필이 길고 배치가 크면 다른 병목이 나타날 수 있습니다.

간단한 벤치 스켈레톤

import time

def bench_prefill_and_decode(run_prefill, run_decode, out_tokens=64):
    t0 = time.perf_counter()
    run_prefill()
    t1 = time.perf_counter()

    for _ in range(out_tokens):
        run_decode()
    t2 = time.perf_counter()

    ttft_ms = (t1 - t0) * 1000
    tpot_ms = ((t2 - t1) * 1000) / out_tokens
    return ttft_ms, tpot_ms

실무에서는 프롬프트 길이를 여러 구간(예: 128, 512, 2048)으로 나누고, 출력 토큰도 고정해 비교해야 회귀를 빨리 잡을 수 있습니다.

“INT4인데 왜 안 빨라지지?” 체크리스트

INT4 적용 후에도 지연이 기대만큼 줄지 않는 흔한 원인은 아래와 같습니다.

1) 실제로는 디퀀타이즈 비용이 지배적

가중치를 INT4로 저장해도 실행 시점에 FP16으로 풀어서 MatMul을 하는 경로라면, 메모리는 줄어도 속도 이득이 제한됩니다.

  • 프로파일에서 DequantizeLinear 또는 유사 노드가 과도하게 보이는지 확인
  • ORT가 INT4 GEMM 커널을 타는지(EP, 빌드 옵션, 모델 패턴) 확인

2) KV 캐시가 병목인데 가중치만 줄였다

LLM 디코드에서 메모리 트래픽은 가중치뿐 아니라 KV 캐시 접근도 큽니다. 프롬프트가 길수록 KV 캐시가 커지고, 토큰마다 읽고 쓰는 비용이 늘어납니다.

  • 컨텍스트 길이 증가에 따라 TPOT이 급격히 나빠지면 KV 캐시 의심
  • 가능하면 KV 캐시 정밀도(예: FP16 유지) 및 저장 방식 최적화 옵션 확인

3) 스레딩/바인딩/런타임 오버헤드가 누적

토큰당 호출 경로가 길면, 커널이 빨라져도 전체가 빨라지지 않습니다.

  • Python 레벨 루프 오버헤드 감소(배치, 바인딩, 최소한의 텐서 변환)
  • ORT 세션 재사용, allocator 재사용
  • 불필요한 로깅/디버깅 제거

4) 모델 그래프가 최적화에 불리한 형태

같은 모델이라도 ONNX 내보내기 방식에 따라 fuse가 깨지고, MatMul 패턴이 분해되어 커널 최적화를 못 타는 경우가 있습니다.

  • export 시 불필요한 Transpose, Reshape가 과도한지 확인
  • ORT optimizer가 fuse할 수 있는 패턴인지 확인

운영 관점: 성능 최적화는 결국 “병목 진단” 문제

양자화는 성능 최적화의 강력한 레버지만, 운영에서는 “어느 구간이 느린지”를 먼저 좁혀야 시행착오가 줄어듭니다. 예를 들어 gRPC 기반 추론 서버에서 메모리나 지연이 튄다면, 애플리케이션 레벨에서도 프로파일링이 필요합니다.

또한 LLM 서빙은 종종 쿠버네티스 환경에서 돌기 때문에, 성능 이슈가 “모델”이 아니라 “노드 리소스 압박, 스로틀링, 재시작 루프”로 보이는 경우도 있습니다. 추론 파드가 갑자기 느려지거나 재시작한다면 아래 같은 진단 글이 도움이 됩니다.

권장 베스트 프랙티스 요약

  • 측정 지표를 쪼개라: TTFT와 TPOT을 분리하고, 프롬프트 길이 구간별로 비교
  • INT4는 “저장 형식”이 아니라 “실행 커널”이 핵심: 실제로 INT4 GEMM 경로를 타는지 프로파일로 확인
  • 런타임 오버헤드를 줄여라: I O binding, 세션 재사용, 스레딩 튜닝
  • KV 캐시를 같이 보라: 컨텍스트가 길어질수록 캐시가 병목이 되기 쉬움
  • 정확도 회귀 테스트를 자동화: 양자화는 품질 저하가 늦게 드러나는 경우가 있어, 대표 프롬프트 세트로 지속 검증

마무리

ONNX Runtime에서 LLM INT4 양자화는 메모리 비용을 크게 낮추고, 특히 디코드 구간의 지연을 줄일 수 있는 실용적인 선택지입니다. 다만 “INT4 적용” 자체보다 중요한 것은, ORT가 실제로 INT4 최적화 커널을 타도록 모델 그래프와 실행 옵션을 맞추고, TTFT와 TPOT을 분리해 병목을 정확히 겨냥하는 것입니다.

다음 단계로는 (1) 프로파일링 결과 기반으로 특정 연산(MatMul, LayerNorm, Attention)의 비중을 확인하고, (2) 입력 길이와 배치 전략을 조정하며, (3) EP별(CPU, CUDA, TensorRT 등) 성능 차이를 동일 조건에서 비교하는 벤치 파이프라인을 만드는 것을 추천합니다.