제5부 추론 최적화와 서빙

5.6 TensorRT 플러그인·커스텀 레이어

5.6.1 이 장에서 배우는 것 / 선수 지식

5.2~5.5에서 TensorRT 엔진을 빌드하고, 동적 shape와 정밀도(INT8·FP8)를 다뤘다. 그런데 빌드가 연산 미지원으로 실패하면? 새 논문의 커스텀 op, TRT가 모르는 활성화, 또는 "이 연산들을 한 커널로 융합하고 싶다"는 성능 요구 — 이때 플러그인(plugin) 으로 직접 만든 CUDA 커널을 TensorRT 엔진에 삽입하는다. 이 장은 그 메커니즘 — 플러그인 인터페이스·등록·직렬화·동적 shape 대응·ONNX 연결·런타임 로딩 — 을 골격 코드와 함께 판다.

플러그인은 양날의 검이다. 미지원 연산을 풀거나 커스텀 융합으로 성능을 얻지만, C++ 인터페이스 구현·직렬화·버전 마이그레이션의 유지보수 비용이 따른다. 그래서 "플러그인을 쓸지" 자체가 먼저 따져야 할 의사결정이다(대안: 모델을 지원 연산으로 재작성·분해).

배우는 것:

• 언제 플러그인이 필요한가 — 미지원 연산·커스텀 op·성능 융합. 그리고 플러그인 전에 시도할 대안(재작성·분해·in-framework 폴백).
• 플러그인의 구조 — TRT가 빌드·실행 시 호출하는 메서드(출력 형상 계산·포맷 지원·enqueue(커널 실행)·직렬화·clone).
• plugin creator·레지스트리 — 이름/네임스페이스·필드(파라미터)로 플러그인을 등록·생성하는 법.
• 동적 shape 플러그인(IPluginV2DynamicExt / IPluginV3 계열) — 형상 표현식·포맷 조합 지원.
• ONNX 연결 — 커스텀 op를 그래프에 넣고 파서가 플러그인에 매핑하게 하기(onnx-graphsurgeon).
• 직렬화·런타임 로딩 — 파라미터를 엔진에 굽고, 추론 시 플러그인 라이브러리를 로드하는 법.

선수 지식: 5.2~5.3(빌드 흐름·동적 shape — 플러그인은 이 위에 삽입된다), 2.1~2.9(CUDA 커널 — 플러그인의 enqueue가 곧 커널 런치), 3.3(cuDNN 융합 — 융합의 대안), 2.12(컴파일·SASS). C++ CUDA 빌드 환경(헤더·nvcc)이 필요하다(0.3/A.1).

정직: 플러그인 API는 TensorRT에서 가장 빠르게 변하는 부분 중 하나다. IPluginV2 → IPluginV2DynamicExt → IPluginV3로 인터페이스가 세대 진화했고(구버전 deprecation), 메서드 시그니처·등록 매크로가 버전마다 다르다. 또 플러그인은 C++로 TRT 헤더에 맞춰 컴파일해야 하므로 이 장의 코드는 복붙 즉시 빌드되는 완본이 아니라 구조를 보여주는 골격이다 — 실제 구현은 당신 TRT 버전의 공식 샘플·헤더를 기준으로 한다("시점 의존, 확인 필요"). Python 플러그인 지원도 버전 의존.

5.6.2 언제 플러그인인가 — 그리고 그 전에

미지원 연산을 만나면 ONNX 파싱이 실패하거나 그 노드에서 멈춘다. 플러그인은 마지막 수단에 가깝다 — 먼저 더 저비용인 대안을 따져라.

견해: 플러그인은 "꼭 필요할 때만". 한 줄 재작성으로 풀리는 것을 플러그인으로 만들면, 그 C++ 코드·직렬화·버전 마이그레이션을 평생 짊어진다. 먼저 (1)(2)(3)을 진지하게 시도하고, 재작성 불가능한 미지원 op이거나 융합 이득이 측정으로 확실할 때만 (4).

5.6.3 플러그인의 구조 — TRT가 부르는 메서드들

플러그인은 TRT가 빌드·실행의 여러 시점에 호출하는 콜백 집합이다. 핵심은 enqueue(실제 CUDA 커널 실행) 이고, 나머지는 TRT가 형상·메모리·직렬화를 다루기 위한 메타 정보다.

enqueue의 시그니처 개요(동적 shape 계열, 버전 의존):

int32_t enqueue(
    const PluginTensorDesc* inputDesc,   // 입력들의 형상·타입·포맷·스케일
    const PluginTensorDesc* outputDesc,  // 출력들의 형상·타입·포맷
    const void* const* inputs,           // 입력 GPU 포인터들
    void* const* outputs,                // 출력 GPU 포인터들
    void* workspace,                     // getWorkspaceSize로 요청한 임시 메모리
    cudaStream_t stream) noexcept;       // 이 스트림에 커널을 런치(비동기, 2.10)

enqueue 안에서 하는 일 = 당신의 CUDA 커널을 그 stream에 런치(2.1~2.3). 즉 플러그인 = "TRT가 형상·메모리·직렬화를 관리해 주는, 그래프에 박힌 커스텀 CUDA 커널".

5.6.4 워크스루 (1) — 플러그인 골격(동적 shape)

간단한 예: 원소별 커스텀 활성화(예: 스케일된 GELU 변형) 플러그인. (구조 골격 — 실제 빌드는 당신 TRT 버전 헤더·샘플 기준.)

5.6.4.1 CUDA 커널

// custom_act_kernel.cu — 실제 계산(보통의 CUDA 커널, 2.1~2.3)
#include <cuda_runtime.h>

__global__ void customActKernel(const float* x, float* y, int n, float alpha) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;   // grid-stride 생략(2.3)
    if (i < n) {
        float v = x[i];
        y[i] = v * 0.5f * (1.f + tanhf(alpha * (v + 0.044715f * v*v*v))); // 예시 식
    }
}

void launchCustomAct(const float* x, float* y, int n, float alpha, cudaStream_t s) {
    int block = 256, grid = (n + block - 1) / block;
    customActKernel<<<grid, block, 0, s>>>(x, y, n, alpha);   // 그 stream에 런치
}

5.6.4.2 플러그인 클래스(핵심 메서드)

// custom_act_plugin.cpp — IPluginV2DynamicExt 계열 골격(시그니처는 버전 확인)
#include <NvInferRuntime.h>
using namespace nvinfer1;
void launchCustomAct(const float*, float*, int, float, cudaStream_t);

class CustomActPlugin : public IPluginV2DynamicExt {
  float mAlpha;                              // 이 플러그인의 파라미터(직렬화 대상)
 public:
  CustomActPlugin(float alpha) : mAlpha(alpha) {}
  // 역직렬화 생성자: 엔진에서 파라미터 복원
  CustomActPlugin(const void* data, size_t len) { mAlpha = *static_cast<const float*>(data); }

  // 출력 형상: 원소별이라 입력과 동일
  DimsExprs getOutputDimensions(int, const DimsExprs* in, int, IExprBuilder&) noexcept override
  { return in[0]; }

  // 지원 포맷/정밀도: 여기선 FP32 linear만(실전은 FP16/INT8도 명시)
  bool supportsFormatCombination(int pos, const PluginTensorDesc* d, int, int) noexcept override
  { return d[pos].type == DataType::kFLOAT && d[pos].format == TensorFormat::kLINEAR; }

  size_t getWorkspaceSize(const DynamicPluginTensorDesc*, int,
                          const DynamicPluginTensorDesc*, int) const noexcept override
  { return 0; }                              // 임시 메모리 불필요

  // ★ 실제 실행: 입력 원소 수를 세어 커널 런치
  int enqueue(const PluginTensorDesc* inDesc, const PluginTensorDesc*,
              const void* const* inputs, void* const* outputs,
              void*, cudaStream_t stream) noexcept override {
    int n = 1; for (int i=0;i<inDesc[0].dims.nbDims;i++) n *= inDesc[0].dims.d[i];
    launchCustomAct(static_cast<const float*>(inputs[0]),
                    static_cast<float*>(outputs[0]), n, mAlpha, stream);
    return 0;                                // 0 = 성공
  }

  // 직렬화: 파라미터를 엔진에 굽기
  size_t getSerializationSize() const noexcept override { return sizeof(float); }
  void serialize(void* buf) const noexcept override { *static_cast<float*>(buf) = mAlpha; }

  DataType getOutputDataType(int, const DataType* in, int) const noexcept override { return in[0]; }
  IPluginV2DynamicExt* clone() const noexcept override { return new CustomActPlugin(mAlpha); }
  const char* getPluginType() const noexcept override { return "CustomAct"; }
  const char* getPluginVersion() const noexcept override { return "1"; }
  int getNbOutputs() const noexcept override { return 1; }
  int initialize() noexcept override { return 0; }
  void terminate() noexcept override {}
  void destroy() noexcept override { delete this; }
  void setPluginNamespace(const char* ns) noexcept override { mNs = ns; }
  const char* getPluginNamespace() const noexcept override { return mNs.c_str(); }
  void configurePlugin(const DynamicPluginTensorDesc*, int,
                       const DynamicPluginTensorDesc*, int) noexcept override {}
 private:
  std::string mNs;
};

위는 구조를 보여주는 축약 골격이다(에러 처리·FP16/INT8 포맷·실제 시그니처 생략). 핵심 흐름: 출력 형상 알려주기 → 포맷 지원 선언 → enqueue에서 커널 런치 → 파라미터 직렬화. 실제 구현은 공식 샘플(plugin 예제)을 베이스로 삼아라.

5.6.5 워크스루 (2) — creator·레지스트리 등록

TRT가 이름으로 플러그인을 찾아 생성하려면 plugin creator를 레지스트리에 등록해야 한다. creator는 "플러그인 타입·버전·필드(파라미터)"를 선언하고, 그 정보로 인스턴스를 만든다.

// custom_act_creator.cpp — creator + 등록(매크로·시그니처 버전 확인)
class CustomActCreator : public IPluginCreator {
  PluginFieldCollection mFC;
  std::vector<PluginField> mFields;          // 파라미터 스키마: "alpha"(float)
 public:
  CustomActCreator() {
    mFields = { PluginField("alpha", nullptr, PluginFieldType::kFLOAT32, 1) };
    mFC.nbFields = mFields.size(); mFC.fields = mFields.data();
  }
  const char* getPluginName() const noexcept override { return "CustomAct"; }   // ONNX op명과 매칭
  const char* getPluginVersion() const noexcept override { return "1"; }
  const PluginFieldCollection* getFieldNames() noexcept override { return &mFC; }

  // 빌드 시 ONNX 노드 속성(alpha)으로 플러그인 생성
  IPluginV2* createPlugin(const char*, const PluginFieldCollection* fc) noexcept override {
    float alpha = 1.0f;
    for (int i=0;i<fc->nbFields;i++)
      if (!strcmp(fc->fields[i].name,"alpha")) alpha = *static_cast<const float*>(fc->fields[i].data);
    return new CustomActPlugin(alpha);
  }
  // 엔진 로드 시 직렬화된 데이터로 복원
  IPluginV2* deserializePlugin(const char*, const void* data, size_t len) noexcept override {
    return new CustomActPlugin(data, len);
  }
  void setPluginNamespace(const char*) noexcept override {}
  const char* getPluginNamespace() const noexcept override { return ""; }
};

REGISTER_TENSORRT_PLUGIN(CustomActCreator);   // ★ 레지스트리에 자동 등록(매크로, 버전 확인)

# 플러그인을 공유 라이브러리로 빌드(TRT·CUDA 헤더·링크 필요)
nvcc -shared -Xcompiler -fPIC \
     custom_act_kernel.cu custom_act_plugin.cpp custom_act_creator.cpp \
     -I/usr/local/cuda/include -I<TRT_ROOT>/include \
     -L<TRT_ROOT>/lib -lnvinfer \
     -o libcustom_act.so

핵심: getPluginName()("CustomAct") = ONNX 그래프의 op 이름으로 매칭된다. 파서가 "CustomAct" 노드를 만나면 이 creator로 플러그인을 만든다. 버전·네임스페이스까지 일치해야 한다.

5.6.6 워크스루 (3) — ONNX 연결·빌드·실행

5.6.6.1 ONNX 그래프에 커스텀 노드 넣기

모델을 ONNX로 export할 때 그 연산이 "CustomAct"라는 op로 표현돼야 파서가 플러그인에 연결한다. 방법:

• torch.onnx의 커스텀 symbolic으로 그 모듈을 CustomAct op로 내보내거나,
• onnx-graphsurgeon으로 export 후 그래프에서 해당 노드를 커스텀 op로 교체.

# graphsurgeon으로 노드를 커스텀 플러그인 op로 치환(개념)
import onnx_graphsurgeon as gs, onnx, numpy as np
graph = gs.import_onnx(onnx.load("model.onnx"))
for node in graph.nodes:
    if node.name == "the_custom_activation":
        node.op = "CustomAct"                 # ★ creator의 getPluginName과 일치
        node.attrs = {"alpha": np.float32(1.0)}  # creator 필드와 일치
graph.cleanup()
onnx.save(gs.export_onnx(graph), "model_plugin.onnx")

5.6.6.2 빌드(플러그인 로드) + 실행

# 빌드 시 플러그인 .so를 로드해 'CustomAct'를 인식시킴
trtexec --onnx=model_plugin.onnx --saveEngine=model.plan \
        --staticPlugins=libcustom_act.so \
        --fp16
# --staticPlugins: 플러그인 .so 로드(플래그명 버전 확인)

런타임(추론)에서도 같은 플러그인 라이브러리를 반드시 로드해야 한다 — 엔진에는 파라미터만 굽혔고, 실제 커널 코드는 .so에 있다:

import ctypes, tensorrt as trt
ctypes.CDLL("libcustom_act.so")               # ★ 추론 프로세스에서 플러그인 로드(필수)
# 또는 표준 플러그인은 trt.init_libnvinfer_plugins(logger, "")
runtime = trt.Runtime(trt.Logger())
engine = runtime.deserialize_cuda_engine(open("model.plan","rb").read())
# 이후 추론은 5.2.5와 동일

함정: 런타임에 플러그인 .so를 로드하지 않으면 deserialize에서 "creator를 못 찾음" 에러. 엔진은 "CustomAct v1을 쓴다"는 참조만 갖고, 구현은 라이브러리에 있다. 빌드 환경과 추론 환경 양쪽에서 .so를 로드해야 한다.

5.6.7 직렬화·동적 shape·정밀도 — 주의점

직렬화. 플러그인의 파라미터(여기선 alpha)는 serialize로 엔진에 구워지고, 로드 시 deserializePlugin으로 복원된다. getSerializationSize와 serialize가 일치해야 하고, 역직렬화 생성자가 같은 레이아웃으로 읽어야 한다 — 어긋나면 조용한 메모리 손상·잘못된 출력. 파라미터를 추가·변경하면 버전 번호를 높여 옛 엔진과 구분한다.

동적 shape. getOutputDimensions는 입력 형상 표현식(DimsExprs)으로 출력 형상을 기호적으로 계산해야 한다(런타임 형상에 대응, 5.3). 원소별이면 입력과 동일하지만, 형상을 바꾸는 연산은 IExprBuilder로 식을 만든다. enqueue는 매 실행의 실제 형상(PluginTensorDesc.dims)을 읽어 커널 크기를 정한다.

정밀도/포맷. supportsFormatCombination이 선언한 포맷만 TRT가 그 플러그인에 보낸다. FP16/INT8을 지원하려면 그 포맷을 선언하고 커널이 그 타입을 처리해야 한다. INT8이면 스케일(PluginTensorDesc.scale)을 받아 처리 — 안 그러면 그 레이어가 정밀도 경계가 되어 주변에 reformat(변환) 오버헤드가 생긴다.

플러그인이 FP32만 지원하는데 주변이 FP16이면, 경계마다 변환이 개입한다:

→ 가능하면 플러그인도 FP16/INT8을 지원해 경계 변환을 없앤다.

5.6.8 중간 점검 (체크포인트)

• [ ] 미지원 연산을 만났을 때 플러그인 전에 시도할 대안(재작성·분해·in-framework)을 대고, 플러그인이 정당한 경우를 구분한다.
• [ ] 플러그인이 TRT가 부르는 콜백 집합이며, enqueue가 곧 CUDA 커널 런치임을 안다.
• [ ] creator·레지스트리로 플러그인을 이름·버전·네임스페이스로 등록·생성하는 흐름을 설명한다.
• [ ] ONNX의 커스텀 op를 플러그인에 연결하고(graphsurgeon), 빌드·런타임 양쪽에서 .so를 로드해야 함을 안다.
• [ ] 직렬화 일치·동적 shape 형상식·포맷/정밀도 지원의 주의점을 제시한다.

증상→의심: deserialize 시 "plugin creator not found" → 런타임에 플러그인 .so 미로드 → 로드. 출력이 쓰레기/크래시 → 직렬화 크기·레이아웃 불일치, 또는 enqueue의 형상 계산 오류 → serialize↔deserialize 점검. 플러그인 주변에 reformat 커널 → 포맷/정밀도 미지원으로 경계 변환 → supportsFormatCombination 확장. 빌드 시 op 인식 못 함 → ONNX op명과 getPluginName 불일치·플러그인 미로드. 버전 업 후 옛 엔진 로드 실패 → 인터페이스/버전 변경 → 엔진 재빌드.

5.6.9 트레이드오프·대안

견해: 재작성 ≫ 분해 ≫ 폴백 ≫ 플러그인 순으로 비용이 커진다. 플러그인의 실질적 비용은 작성이 아니라 TRT 버전이 오를 때마다의 마이그레이션(IPluginV2→V3 같은)이다. 그러니 플러그인은 장기적으로 유지할 가치가 있는 핵심 커스텀 연산에만 투자하고, 일회성·실험은 in-framework 폴백으로 푸는 것이 합리적이다. LLM의 커스텀 어텐션·디코딩은 직접 플러그인보다 TRT-LLM(5.7~)이 이미 최적화해 둔 것을 쓰는 것이 대개 낫다.

비엔비디아 정직: 미지원 op 처리는 다른 런타임도 각자의 커스텀 op 메커니즘을 가진다(ONNX Runtime custom ops 등). TRT 플러그인의 강점은 TRT 그래프 최적화·정밀도·동적 shape와 통합되어 그 GPU에서 가속된다는 점. 단 C++·버전 종속이 부담이니, 정말 필요한지부터 따져라.

5.6.10 흔한 함정

• 플러그인을 너무 일찍. 한 줄 재작성·분해로 풀릴 것을 플러그인화 → 영구 유지보수 부채. 대안을 먼저 소진.
• 런타임에 .so 미로드. 엔진은 참조만 갖는다 — 추론 프로세스에서 플러그인 라이브러리를 로드해야 deserialize 성공.
• 직렬화 불일치. getSerializationSize↔serialize↔역직렬화 생성자가 안 맞으면 조용한 손상. 파라미터 변경 시 버전 상향.
• 동적 shape 형상식 오류. getOutputDimensions를 정적 형상으로 하드코딩하면 동적 입력에서 깨짐. 표현식(DimsExprs) 으로.
• 포맷/정밀도 미지원으로 reformat 폭발. FP32만 지원하면 FP16/INT8 그래프에서 경계마다 변환 커널. 가능한 정밀도를 지원하라.
• 빌드/런타임 TRT 버전 불일치. 플러그인은 헤더 버전에 맞춰 컴파일 — 엔진·런타임과 버전 정합 유지(5.2의 엔진 종속과 동일 원리).
• enqueue에서 동기화·할당. enqueue는 핫패스다 — 안에서 cudaMalloc·동기화하면 성능 붕괴. 메모리는 getWorkspaceSize로 미리 요청, 비동기 stream만 사용.
• 에러를 삼킴. enqueue 실패를 0(성공)으로 반환하면 디버깅 불가. 실제 구현은 에러 코드·로깅을 갖춰야.

5.6.11 연습문제

연습 1 (대안 우선순위). 새 모델에 TRT 미지원인 "커스텀 정규화" op가 있다. 플러그인을 작성하기 전에 검토할 대안 셋을 비용 순으로 들고, 각각이 가능한 조건을 설명하라.

연습 2 (enqueue). 플러그인의 enqueue가 "곧 CUDA 커널 런치"라는 말의 의미를 설명하고, enqueue 안에서 절대 하면 안 되는 일 두 가지와 이유를 제시라.

연습 3 (직렬화·런타임). 플러그인으로 빌드한 엔진을 다른 서버에 배포했더니 deserialize에서 "plugin creator not found"가 난다. 원인과 해결, 그리고 엔진과 플러그인 .so의 역할 분담을 설명하라.

연습 4 (정밀도 경계). FP16 그래프에 FP32만 지원하는 플러그인을 넣었더니 프로파일에 플러그인 앞뒤로 reformat 커널이 보이고 느리다. 무슨 일이며 어떻게 고치나?

연습 5 (유지보수 판단). 팀이 실험용 신규 op를 위해 풀 플러그인을 작성하자고 한다. TRT는 분기마다 버전이 오른다. 이 결정을 어떻게 따져야 하며, 대안 경로는?

5.6.12 요약 · 다음 장

• 플러그인은 TRT 미지원 연산·커스텀 융합을 위해 직접 만든 CUDA 커널을 엔진에 삽입하는 메커니즘이다. 핵심은 enqueue(커널 런치)이고, 나머지는 TRT가 형상·메모리·직렬화를 관리하기 위한 콜백.
• 플러그인 전에 대안(지원 op 재작성·분해·in-framework 폴백)을 소진하라 — 플러그인의 실질적 비용은 버전 마이그레이션 유지보수다.
• 흐름: CUDA 커널 → 플러그인 클래스(형상·포맷·enqueue·직렬화) → creator 등록 → ONNX 커스텀 op 연결 → 빌드·런타임 양쪽 .so 로드. 이름·버전·네임스페이스가 일치해야 한다.
• 주의: 직렬화 일치·동적 shape 형상식·포맷/정밀도 지원(미지원이면 reformat 오버헤드). enqueue 핫패스에서 할당·동기화 금지.
• 정직: 플러그인 API는 TRT에서 가장 변동이 큰 부분(IPluginV2→V3), 코드는 골격이며 당신 버전 공식 샘플·헤더 기준으로 구현. Python 플러그인 지원은 버전 의존.

다음 장: 5.7 TensorRT-LLM I — 아키텍처·KV 캐시. 여기서 일반 TensorRT를 떠나 LLM 특화 런타임으로 들어간다. 자기회귀 디코딩(5.1.6)을 위한 KV 캐시가 무엇이고 왜 추론 메모리·처리량의 중심인지, TensorRT-LLM이 일반 TRT와 무엇이 다르고 어떻게 LLM을 빌드·실행하는지를 본다.