본문 바로가기

dev_AI_framework

cuBLASLt / CUTLASS 의 에필로그 epilogue 까지 흡수한다. - GEMM 연산 이후의 연산들 (bias add, activation, scaling 등 ) 의 커널 통합 ( 호스트 디스패치 + 커널 특화 방식, 나는 CUTLASS 가 맞을 듯??)

1. GEMM 연산의 구조

보통 C = A @ B + bias 같은 연산은 이렇게 처리

  1. 메인 커널 (Main GEMM) : 행렬곱 D = A @ B  계산
  2. 후처리 (Post-proccessing) : bias add, activation function, scaling, cilp, etc... 이런 것들을 별도 커널로 다시 실행

이 경우 launch 오버헤드 발생, 메모리 read & write 과정이 추가되어 성능 감소

 

2. Epilogue 흡수

cuBLASLt, CUTLASS 는 GEMM 내부 구조를 크게 두 부분으로 나눈다. 

  • Mainloop : 행렬곱의 핵심 루프 (tile 단위 accumulate)
  • Epilogue : 결과를 global memory 에 저장하기 직전 단계

여기서 Epilogue Visitor 라는 확장 포이트가 있어서, accumulator register 안에 있는 결과를 메모리에 쓰기 전에 bias 를 더하거나activation 의 적용, scaling 을 할 수 있다.

 

3. 장점

  • 메모리 왕복 제거 : bias/activation 을 위해 다시 global memory load/store 필요 없음
  • 커널 런치 오버헤드 감소 : 하나의 GEMM 호출로 끝남
  • 성능 최적화 : bandwidth 절약 + pipeline 효율상승

 

 

4. 예시

  • cuBLASLt 에서는 cublasLtMatmulDescSetAttribute 로 epilogue 동작 지정
  • CUTLASS 에서는 EpilogueVisitor 클래스를 구현해서, acc = f(bias, activation(acc)) 이런 식으로 커스터마이즈 
cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS_ACT;  
cublasLtMatmulDescSetAttribute(matmulDesc,
    CUBLASLT_MATMUL_DESC_EPILOGUE,
    &epilogue,
    sizeof(epilogue));

 

  • cuBLASLt : 에필로그에서 뭘 할지를 런타임 옵션으로 고르는 방식, 별도 커널을 여러 개 구현하는 것이 아닌, 미리  정해진 몇 가지 조합을 지원하고, 그 중 하나를 호스트 코드에서 설정하면 내부가 알아서 해당 변형의 커널로 디스패치
  • CUTALSS 
    • 컴파일타임 특화(권장, 최고 성능) - 에필로그 연산 (Bias + ReLU, Bias+GELU...) 별로 템플릿 인스턴스를 만들어 두고, 호스트에서 switch 로 선택해서 호출, 커널은 여러 개지만, 각 커널은 분기 없이 깔끔
    • 런타임 분기 ( 유연하지만 성능 손해 가능 ) - 하나의 커널 에필로그 안에서 switch(op) 로 여러 연산 처리, 코드 부피/레지스터 압박과 워프 분기 리스트 존재 

 

왜 호스트 디스패치 + 커널 특화가 일반적인가?

  • 에필로그는 레지스터에 있는 누산값(accumulator) 을 메모리 에 쓰기 직전에 실행, 여기에서 복잡한 분기/함수 포인터를 쓰면 인라이닝이 깨지고 레지스터 압박이 커져 성능이 확 떨어진다.
  • 그래서 여러 조합을 미리 템플릿으로 만든 커널을 등록해두고, 호스트에서 enum/switch 로 선택 -> 해당 톡화 커널 호출이 보편적 패턴

 

특수한 활성화/커스텀 연산이 필요할 경우, CUTLASS 의 Epilogue Visitor/Functor 로 구현 -> 해당 조합을 템플릿 인스턴스로 추가 -> 호스트 디스패치 ( 런타임 단일 커널 내 분기보다는, 조합별 커널 특화를 권장 )