GPU용 TritonSigmoid 공개

TritonSigmoid를 오픈소스로 공개했다. 싱글셀 foundation models에서 각 세포를 gene sequence로 표현할 때 쓰는 padding-aware sigmoid attention GPU 커널로, 200~16,000+ 토큰의 가변 길이 입력을 빈 패딩 계산 없이 다룬다.

Softmax가 토큰 간 경쟁을 강제하는 반면, sigmoid는 여러 유전자와 전사인자에 동시에 강한 attention을 허용한다.

주요 결과는 다음과 같다.

• H100에서 최대 515 TFLOPS를 기록했다.
• FlashAttention-2(361), **FlashSigmoid(440)**보다 높은 처리량을 보였다.
• 6개 held-out 데이터셋에서 softmax attention보다 낮은 validation loss를 냈다.
• cell-type separation이 25% 개선됐다.
• softmax가 발산하던 조건에서도 학습이 안정적으로 유지됐다.

논문은 arXiv에, 구현은 GitHub에 공개됐다.