mirror of
https://github.com/huggingface/candle.git
synced 2025-06-15 18:28:24 +00:00
59 lines
1.1 KiB
C++
59 lines
1.1 KiB
C++
#ifndef _GPU_OPS_KERNELS_H_
|
|
#define _GPU_OPS_KERNELS_H_
|
|
|
|
#include <cuda_runtime_api.h>
|
|
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
|
|
#include<stdlib.h>
|
|
#include<stdint.h>
|
|
|
|
namespace gpu_ops {
|
|
|
|
struct MHAParams {
|
|
uint32_t q_batch_stride;
|
|
uint32_t k_batch_stride;
|
|
uint32_t v_batch_stride;
|
|
uint32_t o_batch_stride;
|
|
|
|
uint32_t q_row_stride;
|
|
uint32_t k_row_stride;
|
|
uint32_t v_row_stride;
|
|
uint32_t o_row_stride;
|
|
|
|
uint32_t q_head_stride;
|
|
uint32_t k_head_stride;
|
|
uint32_t v_head_stride;
|
|
uint32_t o_head_stride;
|
|
|
|
uint32_t b;
|
|
uint32_t h;
|
|
uint32_t h_k;
|
|
uint32_t d;
|
|
uint32_t d_rounded;
|
|
float softmax_scale;
|
|
float softcap;
|
|
|
|
uint32_t seqlen_q;
|
|
uint32_t seqlen_k;
|
|
uint32_t seqlen_q_rounded;
|
|
uint32_t seqlen_k_rounded;
|
|
|
|
int window_size_left;
|
|
int window_size_right;
|
|
|
|
int is_causal;
|
|
int is_bf16;
|
|
};
|
|
|
|
void run_mha_fwd_j(cudaStream_t stream, void **buffers,
|
|
const char *opaque,
|
|
std::size_t opaque_len);
|
|
void run_mha_bwd_j(cudaStream_t stream, void **buffers,
|
|
const char *opaque,
|
|
std::size_t opaque_len);
|
|
}
|
|
|
|
#endif
|