wehub-resource-sync
132 lines
2.8 KiB
Plaintext
132 lines
2.8 KiB
Plaintext
#ifdef __HIP_PLATFORM_AMD__
|
|
#include "ds4_rocm.h"
|
|
#include <hipblaslt/hipblaslt.h>
|
|
|
|
#define FULL_WARP_MASK 0xFFFFFFFFFFFFFFFFULL
|
|
#define MASK_T uint64_t
|
|
#define DS4_GPU_BACKEND_NAME "ROCm"
|
|
#define DS4_GPU_LOG_PREFIX "ds4: ROCm "
|
|
#define DS4_GPU_BLAS_NAME "hipBLAS"
|
|
#else
|
|
#include <cuda_runtime.h>
|
|
#include <cuda_fp16.h>
|
|
#include <mma.h>
|
|
#include <cublas_v2.h>
|
|
#include <cub/block/block_radix_sort.cuh>
|
|
|
|
#define FULL_WARP_MASK 0xFFFFFFFFu
|
|
#define MASK_T uint32_t
|
|
#define DS4_GPU_BACKEND_NAME "CUDA"
|
|
#define DS4_GPU_LOG_PREFIX "ds4: CUDA "
|
|
#define DS4_GPU_BLAS_NAME "cuBLAS"
|
|
#endif
|
|
|
|
#include <stdint.h>
|
|
#include <ctype.h>
|
|
#include <errno.h>
|
|
#include <limits.h>
|
|
#include <math.h>
|
|
#include <fcntl.h>
|
|
#include <pthread.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <sys/stat.h>
|
|
#include <time.h>
|
|
#include <unistd.h>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
#include "ds4_gpu.h"
|
|
|
|
#ifndef M_PI
|
|
#define M_PI 3.14159265358979323846
|
|
#endif
|
|
|
|
#define CUDA_QK_K 256
|
|
#define DS4_ROCM_UNUSED __attribute__((unused))
|
|
|
|
enum {
|
|
/* attention_decode_mixed_kernel stores raw-window scores plus visible
|
|
* compressed scores in shared memory. The host routes larger unmasked
|
|
* decode calls to the online attention kernel so this fixed buffer never
|
|
* becomes an out-of-bounds write at long context. */
|
|
DS4_ROCM_ATTENTION_SCORE_CAP = 8192u,
|
|
DS4_ROCM_ATTENTION_RAW_SCORE_CAP = 256u,
|
|
DS4_ROCM_TOPK_MERGE_GROUP = 8u
|
|
};
|
|
|
|
struct ds4_gpu_tensor {
|
|
void *ptr;
|
|
uint64_t bytes;
|
|
int owner;
|
|
};
|
|
|
|
typedef struct {
|
|
uint8_t scales[CUDA_QK_K / 16];
|
|
uint8_t qs[CUDA_QK_K / 4];
|
|
uint16_t d;
|
|
uint16_t dmin;
|
|
} cuda_block_q2_K;
|
|
|
|
typedef struct {
|
|
uint16_t d;
|
|
uint16_t dmin;
|
|
uint8_t scales[12];
|
|
uint8_t qs[CUDA_QK_K / 2];
|
|
} cuda_block_q4_K;
|
|
|
|
typedef struct {
|
|
float d;
|
|
int8_t qs[CUDA_QK_K];
|
|
int16_t bsums[CUDA_QK_K / 16];
|
|
} cuda_block_q8_K;
|
|
|
|
typedef struct {
|
|
uint16_t d;
|
|
uint16_t qs[CUDA_QK_K / 8];
|
|
} cuda_block_iq2_xxs;
|
|
|
|
#include "ds4_iq2_tables_cuda.inc"
|
|
|
|
#include "rocm/ds4_rocm_runtime.cuh"
|
|
|
|
#include "rocm/ds4_rocm_common.cuh"
|
|
|
|
#include "rocm/ds4_rocm_q8.cuh"
|
|
|
|
#include "rocm/ds4_rocm_norm_rope.cuh"
|
|
|
|
#include "rocm/ds4_rocm_fp8_kv.cuh"
|
|
|
|
#include "rocm/ds4_rocm_attention.cuh"
|
|
|
|
#include "rocm/ds4_rocm_hc.cuh"
|
|
|
|
#include "rocm/ds4_rocm_output.cuh"
|
|
|
|
#include "rocm/ds4_rocm_indexer.cuh"
|
|
|
|
#include "rocm/ds4_rocm_embedding_launch.cuh"
|
|
|
|
#include "rocm/ds4_rocm_matmul.cuh"
|
|
|
|
#include "rocm/ds4_rocm_fp8_kv_launch.cuh"
|
|
|
|
#include "rocm/ds4_rocm_compressor.cuh"
|
|
|
|
#include "rocm/ds4_rocm_attention_launch.cuh"
|
|
|
|
#include "rocm/ds4_rocm_shared_expert.cuh"
|
|
|
|
#include "rocm/ds4_rocm_misc_launch.cuh"
|
|
#include "rocm/ds4_rocm_router.cuh"
|
|
|
|
#include "rocm/ds4_rocm_moe.cuh"
|
|
|
|
#include "rocm/ds4_rocm_moe_launch.cuh"
|
|
|
|
#include "rocm/ds4_rocm_hc_output_launch.cuh"
|
|
|
|
#include "rocm/ds4_rocm_current_api_compat.cuh"
|