File size: 1,422 Bytes
462dacf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
#ifndef _cuda_buffers_cuh
#define _cuda_buffers_cuh
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cstdint>
#include <cstdio>
const int CUDA_MAX_DEVICES = 16;
// #ifndef _cuda_buffers_cu
// extern __constant__ half2 q4_table[16][256];
// #endif
class CudaBuffers
{
public:
int device;
half* temp_state; // [max_hidden_rows * intermediate_size]
int temp_state_size;
half* temp_mlp; // [hidden_dim * intermediate_size]
float* temp_zeros_float; // [max_hidden_rows]
half* temp_dq; // size of largest quant tensor * 8
int current_zeros_float;
int max_zeros_float;
cudaStream_t alt_stream_1;
cudaStream_t alt_stream_2;
cudaStream_t alt_stream_3;
cudaEvent_t alt_stream_1_done;
cudaEvent_t alt_stream_2_done;
cudaEvent_t alt_stream_3_done;
CudaBuffers
(
int _device,
half* _temp_state,
int _temp_state_size,
half* _temp_mlp,
float* _temp_zeros_float,
half* _temp_dq,
int _max_zeros_float
);
~CudaBuffers();
float* get_zeros_float(const int num_zeros);
};
CudaBuffers* get_buffers(const int device_index);
void prepare_buffers_cuda
(
int _device,
half* _temp_state,
int _temp_state_size,
half* _temp_mlp,
float* _temp_zeros_float,
half* _temp_dq,
int _max_zeros_float
);
void cleanup_buffers_cuda();
#endif |