정책화된 커널 내부 에필로그 호출 방식

1. 개념 요약

정책화된 커널 내부 에필로그 호출이란 GEMM, Conv 같은 메인 연산 커널 내부에서, 연산 결과를 최종 저장하기 전에 후처리를 정책으로 분리하고, 이 정책을 템플릿 인라인 함수 형태로 호출하는 기법

 

커널 안에서 이렇게 일어남

float acc = accumulate(A, B);
EP::apply(D, C, Z, p, m, n, acc, bias_m, bias_n);

 

• 여기서 EP 는 Epilogue<ActKind, BiasMode, HasC, SaveZ> 같은 정책 타입
• apply()는 장치 함수 (__device__ __forceinline__ ) - GPU 함수??
• 실제 CUDA 커널은 단 하나만 실행된다. ( EP::apply 가 별도 커널 런치가 아니라 인라인됨)

 

2. 구조적 철학

이 방식을 이해하기 위해 커널을 두 부분으로 나눠본다.

Main Loop (Core)	Accumulate: acc[i][j] += A_tile * B_tile	GEMM/Conv 계산 본체
Epilogue (Postprocess)	Scale + Add + Bias + Activation + Store	결과 후처리 + 출력

이 두 영역을 분리하고, Epilogue 부분만 정책화시킨 것

그럼 커널 내부 구조는 아래와 같이 됌

template<int BM, int BN, int BK, ActKind AK, BiasMode BMmode, bool HasC, bool SaveZ>
__global__ void gemm_kernel(GemmParams p) {
    // [1] Accumulate phase
    float acc[THR_M][THR_N] = {0};
    ...
    acc[i][j] = fmaf(a, b, acc[i][j]);

    // [2] Epilogue phase
    using EP = Epilogue<AK, BMmode, HasC, SaveZ>;
    EP::apply(D, C, Z, p, m, n, acc[i][j], bias_m, bias_n);
}

• 커널 로직은 항상 같고, 
• EP::apply() 가 어떤 일을 하느냐에 따라 후처리가 달라진다.

 

3. Policy 의 역할

정책화는 결국 커널의 에필로그를 교체 가능한 플러그인으로 만드는 것,

각 정책은 템플릿으로 정의됨

template<ActKind AK, BiasMode BMmode, bool HasC, bool SaveZ>
struct Epilogue {
  __device__ __forceinline__
  static void apply(float* D, int ldd,
                    const float* C, int ldc,
                    float* Z, int ldZ,
                    const GemmParams& p,
                    int m, int n,
                    float acc, float bias_m, float bias_n) 
  {
    float pre = p.alpha * acc;

    if constexpr (HasC)
      pre = fmaf(p.beta, C[m * ldc + n], pre);

    if constexpr (BMmode == BiasMode::PerN)
      pre += bias_n;
    else if constexpr (BMmode == BiasMode::PerM)
      pre += bias_m;

    if constexpr (SaveZ)
      Z[m * ldZ + n] = pre;

    float out = apply_act_static<AK>(pre, p.leaky_slope);
    D[m * ldd + n] = out;
  }
};

여기서 constexpr if 로 각 조합을 컴파일에서 결정하기 때문에 런타임 분기가 사라짐

 

4 .장점

병합 위치	커널 밖 (다단계)	커널 안 (단일 실행)
메모리 경로	acc → global → 다시 load → act → store	acc (register) → act → store
런치 횟수	여러 개 (GEMM, bias, act 별도)	1개
분기 오버헤드	런타임 if/else 존재	컴파일타임 constexpr로 제거
Graph Capture 호환성	여러 커널 캡처 불안정	단일 커널이라 안전
성능	글로벌 왕복 + 분기 → 손실	inlined path로 최대 효율
확장성	구현은 단순하나 중복 많음	재사용성 높고 정책만 추가하면 확장

 

5. 작동 순서 예시

1. 런처가 런타임에 따라 컴파일타임 조합을 선택,
2. 이 런처가 특정 커널 실행
3. 커널 안에서 계산을 마친 뒤 EP::apply 호출
4. 이 apply 는 인라인되어 PTX 코드로 직접 들어가며, 별도 런치 없이 활성화, 바이어스, C-add 등을 처리하고 결과 저장

 

6. 중요한 설계 포인트

__device__ __forceinline__

에필로그가 완전히 인라인되어 레지스터 상에서 실행되도록 보장

 

템플릿 정책

act, bias 모드, saveZ 여부 등은 컴파일타임 상수화

 

런처에서 분기 

런처가 런타임 분기를 컴파일타임으로 내림

 

Acc - Epilogue 직접 연결

acc 의 레지스터 값을 그대로 후처리에 사용

 

Reusability 

GEMM/Conv/RNN 이 모두 동일한 Epilogue<> 인터페이스 사용 가능

 

Z-stash / SaveZ 

활성화 전 값을 저장할 수 있는 확장 포인트

 

AMP/Dropout 등 확장

추가 정책 인자만 붙이면 동일 구조로 확장 가능

 

7. 실제 예시로 본 인라인 구조 ( PTX 관점 )

EP::apply() 가 인랑니되면 아래처럼 PTX 가 구성

mul.f32     %f3, %facc, %falpha;      // pre = alpha * acc
fma.rn.f32  %f3, %fbeta, %fC, %f3;    // if HasC
add.f32     %f3, %f3, %fbias;         // add bias
max.f32     %f3, %f3, 0f00000000;     // ReLU
st.global.f32 [D], %f3;               // store result

즉, 커널 내에서 바로 끝남, EPilogue 는 코드 조각 그 자체로 삽입됨

 

8. 유지보수 및 확장

해당 구조의 강점은, 하나의 설계 패턴으로 

모든 모듈을 덮을 수 있다는 점

각 커널은 core loop 만 다르고, 마지막 한 줄은 동일하게 EP::apply()