본문 바로가기

dev_AI_framework

📄 신규 기능 추가 과정 문서: FP16 TensorCore GEMM (cuBLASLt)

1. 목표

  • 기존 graph_executor_v2는 FP32 GEMM 스텁만 지원.
  • 새로 FP16 Tensor Core 최적화 GEMM을 cuBLASLt 기반으로 추가.
  • Bias + Activation(ReLU)까지 지원.
  • Python에서 동일 인터페이스(launch_kernel)로 실행 가능.

 

2. 수정/추가된 파일

2.1 CMakeLists.txt

  • CUDA::cublasLt를 링크에 추가 → cuBLASLt API 사용 가능.

 

2.2 src/launch_table.cpp

  • 커널 테이블에 FP16 버전 추가.
  • Capability 점수를 FP32보다 높게 설정해 자동 선택 우선권 부여.

 

2.3 src/my_kernels.cu

(a) 공통 구조체

  • Python에서 넘겨주는 Params 구조체 (M,N,K, has_bias, act).

(b) FP32 버전 유지

  • 단순 CUDA 커널 (for-loop GEMM).

(c) FP16 + cuBLASLt 신규 구현

  • cuBLASLt로 GEMM 실행: A/B/D는 FP16, accumulate는 FP32.
  • 에필로그(BIAS/ACT)는 사용하지 않고, 대신 후처리 CUDA 커널로 처리.
    • 이유: cuBLASLt는 bias 벡터 길이를 M(행) 기준으로 기대하는데, DL은 N(열) 기준이므로 mismatch 발생.
    • 따라서 GEMM만 하고, bias+ReLU는 별도 add_bias_relu_fp16 커널에서 수행.
  • ge2_launch_gemm_bias_act_tc_f16:
    • GEMM 실행(cuBLASLt).
    • 실행 성공 후 add_bias_relu_fp16 호출.
    • bias: FP32 벡터(N), ReLU: 옵션.

 

2.6 smoke_test_v2_f16.py (신규 테스트)

  • CuPy 기반 검증 스크립트 작성.
  • A/B: FP16, Bias: FP32, C_out: FP16.
  • Python에서 참값(FP32 accumulate → FP16 cast)을 계산해 오차 확인.

 

 

4. 기능 추가 절차 요약 (패턴)

앞으로 다른 기능(예: GELU, LayerNorm, Conv2D)을 추가할 때는 같은 절차를 따르면 됨.

  1. CMakeLists.txt
    • 필요한 CUDA 라이브러리 링크 추가.
  2. my_kernels.cu
    • 커널 함수 구현 (ge2_launch_xxx).
    • 필요하면 후처리 CUDA 커널 작성.
  3. launch_table.cpp
    • 커널 테이블 + capability 테이블에 신규 함수 등록.
  4. bindings_min_api.cpp
    • 별도 수정 불필요 (자동 반영).
  5. 테스트 스크립트 작성
    • CuPy/Torch 등으로 참값 계산 → 오차 확인.

 

5. 이번 변경의 의의

  • FP32 스텁 → FP16 Tensor Core 고성능 GEMM 확장.
  • Bias+ReLU 지원으로 실제 DL 연산에 더 근접.
  • cuBLASLt의 Epilogue 한계를 회피하고 안정성 확보.
  • 프레임워크에서 자동 커널 선택 (capability+selector) 시 FP16 경로를 우선 사용 가능.