#pragma once #ifdef USE_ROCM #include #endif #ifdef USE_ROCM struct Utils { static __host__ int get_warp_size() { static bool is_cached = false; static int result; if (!is_cached) { int device_id; cudaDeviceProp deviceProp; cudaGetDevice(&device_id); cudaGetDeviceProperties(&deviceProp, device_id); result = deviceProp.warpSize; is_cached = true; } return result; } static __device__ constexpr int get_warp_size() { #ifdef __GFX9__ return 64; #else return 32; #endif } }; #define WARP_SIZE Utils::get_warp_size() #else #define WARP_SIZE 32 #endif #ifndef USE_ROCM #define VLLM_LDG(arg) __ldg(arg) #else #define VLLM_LDG(arg) *(arg) #endif #ifndef USE_ROCM #define VLLM_SHFL_XOR_SYNC(var, lane_mask) \ __shfl_xor_sync(uint32_t(-1), var, lane_mask) #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \ __shfl_xor_sync(uint32_t(-1), var, lane_mask, width) #else #define VLLM_SHFL_XOR_SYNC(var, lane_mask) __shfl_xor(var, lane_mask) #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \ __shfl_xor(var, lane_mask, width) #endif #ifndef USE_ROCM #define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane) #else #define VLLM_SHFL_SYNC(var, src_lane) __shfl(var, src_lane) #endif #ifndef USE_ROCM #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) \ __shfl_down_sync(uint32_t(-1), var, lane_delta) #else #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) __shfl_down(var, lane_delta) #endif #ifndef USE_ROCM #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \ cudaFuncSetAttribute(FUNC, cudaFuncAttributeMaxDynamicSharedMemorySize, VAL) #else #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \ hipFuncSetAttribute(FUNC, hipFuncAttributeMaxDynamicSharedMemorySize, VAL) #endif