Fp8 커널 이름에 "cutlass"를 넣으면 약 100TFLOPS 빨라짐

1 month ago 13

요약: Persistent Attention 커널의 리팩토링 및 "cutlass" 네이밍 효과

이 Pull Request는 Triton의 attention 커널을 persistent attention 방식으로 리팩토링한 내용임
주요 목적은 낮은 context 길이 구간에서 성능 최적화 효과를 얻기 위함임
하지만, fp16 유형에서는 context 크기가 커질수록 softmax 부분의 ptxas 명령 스케줄 이슈로 인해 성능 저하 현상이 나타남

fp8 모델을 사용할 때, 커널 이름에 "cutlass"를 포함시키면 약 100TFLOPS까지 성능 향상이 측정됨
이는 NVIDIA의 ptxas(PTX 어셈블러) 가 내부적으로 커널 이름에 "cutlass"가 포함되어 있을 때 특수 최적화를 적용하기 때문임
실제 코드에서 dtype이 float8e5이면 커널 이름을 "cutlass_gluon_attention" 등으로 할당함
ptxas 디스어셈블 분석 결과, 실제로 내부 코드에서 strstr(kernel_name, "cutlass") 조건문이 존재함

리팩토링 전후 benchmark 데이터에서, persistent attention 적용 전보다 D=64에서 상대적 성능 하락이 관측됨
- 이는 softmax 파트의 명령 스케줄 문제에 기인함
하지만 fp8 타입과 "cutlass" 네이밍을 조합할 경우, 특정 context와 크기에서는 현저히 높은 처리량 기록임
대표적인 결과 요약:
- Attention Z=4, H=32, D=64, causal=False:
  - persistent 적용 전: triton-fp16 약 383, triton-fp8 약 413, cudnn-fp16 약 565
  - persistent 적용 후: triton-fp16 약 360, triton-fp8 약 370
- Attention Z=4, H=32, D=128, causal=True:
  - persistent 적용 전: triton-fp16 약 312, triton-fp8 약 345, cudnn-fp16 약 553
  - persistent 적용 후: triton-fp16 약 356, triton-fp8 약 351

"cutlass"라는 문자열이 커널 이름에 포함될 경우, NVIDIA ptxas에서 실험적 최적화 루틴이 활성화됨이 커뮤니티 분석에서 확인됨
ptxas에 네이밍 매칭용 하드코딩 로직(strstr(kernel_name, "cutlass"))이 있음
이 트릭은 aggressive, experimental 최적화로, 커널의 정확도 변화 여부는 하드웨어 및 상황에 따라 달라질 수 있음

커널의 이름만으로도 하드웨어 레벨에서 실질적인 성능 변동이 발생하는 점이 매우 특이 현상임
NVIDIA 소프트웨어 스택(pxtas)에 숨은 최적화 trigger가 존재한다는 점에서, AI 프레임워크 개발 시 네이밍 규칙도 성능에 영향 줄 수 있음을 시사함
실무적으로는 트릭 사용 시 재현성 및 안정성을 반드시 테스트해야 하며, 하드웨어 벤더의 최적화 정책을 예의주시할 필요가 있음