Attention/Transformer Layer 를 구현하려면 - 기저 연산들의 구현 필요

기저 연산(primitive ops)”들을 먼저 제대로 만들고, 그걸 Attention/Transformer 레이어가 조합해서 그래프(IR)로 컴파일 → CUDA 실행 플랜으로 내리면 된다.

 

1) 최소 기저 연산(커널/연산자) 체크리스트

Attention/Transformer에 필요한 “필수”와 “있으면 좋은” 항목을 분리했어요.

필수 (Forward/Backward 지원)

• GEMM / Batched GEMM: Q= XWq, K= XWk, V= XWv, O = (softmax(QK^T/√Dh)V)Wo
  • fp16/bf16 입력 + fp32 누적(Tensor Core) 권장, cublasLt 사용 시 bias+activation 일부 fuse 가능
• Softmax(마스킹 포함): additive mask(0 또는 -inf) + padding mask(bool→additive)
  • 안정화: x - max(x) 후 exp/sum (backward도 안정형)
• LayerNorm / (또는 RMSNorm): pre-norm 권장 (학습 안정+딥스택 유리)
• Dropout: Philox 기반 RNG, in-place 가능(마스크 저장)
• Elementwise: bias_add, add(residual), mul, gelu(또는 silu), relu
• Reshape/Transpose/Permute: (B,T,C) ↔ (B,H,T,Dh) 변환
• Concat/Slice: KV 캐시 확장 시 필수
• Embedding / Positional Encoding: lookup + (sin/cos) or RoPE(회전)
• Mask 유틸: causal(상삼각 -inf), padding(bool→additive)

 

7) 요약 (당신의 질문에 딱 대답)

 

먼저 GEMM/Softmax/LayerNorm/Dropout/Reshape/Mask 같은 기저 연산 레이어들을 견고하게 만들고, Attention/Transformer 레이어에서 그들을 그래프로 조합하세요. 그런 다음 그래프 컴파일러가 연산자 fusion(예: QKV, Add+Dropout+LN), FlashAttention 대체, 메모리 플래닝을 수행해서 CUDA 커널/라이브러리 호출로 내리시면 됩니다.