mirror of
https://github.com/karpathy/llm.c.git
synced 2026-07-26 20:15:08 -04:00
21 lines
No EOL
757 B
C
21 lines
No EOL
757 B
C
/*
|
|
cuDNN (flash) attention
|
|
*/
|
|
#ifndef CUDNN_ATT_H
|
|
#define CUDNN_ATT_H
|
|
|
|
#include "cuda_common.h"
|
|
|
|
// forward declarations of functions defined in cudnn_att.cpp
|
|
void create_cudnn();
|
|
void destroy_cudnn();
|
|
void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS)
|
|
float* stats, // output for backward pass: (B, NH, T)
|
|
floatX* inp, // input: (B, T, 3, NH, HS) QKV
|
|
int B, int T, int NH, int C);
|
|
|
|
void attention_backward_cudnn(floatX* dqkvr, // output
|
|
floatX* dout, floatX* qkvr, floatX* o, float* stats, // inputs
|
|
int B, int T, int NH, int C);
|
|
|
|
#endif // CUDNN_ATT_H
|