char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
074944998
char/aliceai-t5-moe
1#include "argsort.cuh" 2 3#ifdef GGML_CUDA_USE_CUB 4# include <cub/cub.cuh> 5# if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 1) 6# define STRIDED_ITERATOR_AVAILABLE 7# include <cuda/iterator> 8# endif 9using namespace cub; 10#endif // GGML_CUDA_USE_CUB 11 12static __global__ void init_indices(int * indices, const int ncols, const int nrows) { 13 const int col = blockIdx.x * blockDim.x + threadIdx.x; 14 const int row = blockIdx.y; 15 16 if (col < ncols && row < nrows) { 17 indices[row * ncols + col] = col; 18 } 19} 20 21#ifndef STRIDED_ITERATOR_AVAILABLE 22static __global__ void init_offsets(int * offsets, const int ncols, const int nrows) { 23 const int idx = blockIdx.x * blockDim.x + threadIdx.x; 24 if (idx <= nrows) { 25 offsets[idx] = idx * ncols; 26 } 27} 28#endif // STRIDED_ITERATOR_AVAILABLE 29 30#ifdef GGML_CUDA_USE_CUB 31 32// returns the suggested maximum number of rows to process during one argsort_f32_i32_cuda_cub() call 33int argsort_f32_i32_cuda_cub_chunk_nrows(const size_t nb01, const int64_t nrows) { 34 // perform argsort in chunks up to approximately this size (currently 64MB) 35 // to avoid excessive temporary buffers memory usage 36 const int chunk_bytes = 1 << 26; 37 38 // calculate how many rows will fit in one chunk (must be at least one) 39 const int chunk_nrows = std::max((int) (chunk_bytes / nb01), 1); 40 41 // limit the resulting amount to total nrows 42 return std::min((int64_t) chunk_nrows, nrows); 43} 44 45void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, 46 const float * x, 47 int * dst, 48 const int ncols, 49 const int nrows, 50 ggml_sort_order order, 51 cudaStream_t stream) { 52 ggml_cuda_pool_alloc<int> temp_indices_alloc(pool, ncols * nrows); 53 ggml_cuda_pool_alloc<float> temp_keys_alloc(pool, ncols * nrows); 54 55 int * temp_indices = temp_indices_alloc.get(); 56 float * temp_keys = temp_keys_alloc.get(); 57 58 static const int block_size = 256; 59 const dim3 grid_size((ncols + block_size - 1) / block_size, nrows); 60 init_indices<<<grid_size, block_size, 0, stream>>>(temp_indices, ncols, nrows); 61 62#ifdef STRIDED_ITERATOR_AVAILABLE 63 auto offset_iterator = cuda::make_strided_iterator(cuda::make_counting_iterator(0), ncols); 64#else 65 // offset_iterator needs to populate nrows + 1 elements, so we also have to ceildiv nrows + 1 by block_size 66 const int nrows_offset = nrows + 1; 67 ggml_cuda_pool_alloc<int> offsets_alloc(pool, nrows_offset); 68 int * offset_iterator = offsets_alloc.get(); 69 const dim3 offset_grid((nrows_offset + block_size - 1) / block_size); 70 init_offsets<<<offset_grid, block_size, 0, stream>>>(offset_iterator, ncols, nrows); 71#endif 72 CUDA_CHECK(cudaMemcpyAsync(temp_keys, x, ncols * nrows * sizeof(float), cudaMemcpyDeviceToDevice, stream)); 73 74 size_t temp_storage_bytes = 0; 75 76 bool is_capturing = false; 77#ifdef USE_CUDA_GRAPH 78 // Currently (confirmed for CCCL <= 3.2) DeviceSegmentedSort does not support stream capture, while DeviceSegmentedRadixSort does. 79 // See https://github.com/NVIDIA/cccl/issues/5661#issuecomment-3229037149 80 // TODO: constrain this to the CCCL versions that have this issue once it's resolved in a future CCCL release. 81 cudaStreamCaptureStatus capture_status; 82 CUDA_CHECK(cudaStreamIsCapturing(stream, &capture_status)); 83 is_capturing = (capture_status != cudaStreamCaptureStatusNone); 84#endif // USE_CUDA_GRAPH 85 86 if (order == GGML_SORT_ORDER_ASC) { 87 if (nrows == 1) { 88 CUDA_CHECK(DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, temp_keys, temp_keys, // keys (in-place) 89 temp_indices, dst, // values (indices) 90 ncols, 0, sizeof(float) * 8, stream)); 91 } else if (is_capturing) { 92 CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs( 93 nullptr, temp_storage_bytes, temp_keys, temp_keys, // keys (in-place) 94 temp_indices, dst, // values (indices) 95 ncols * nrows, nrows, // num items, num segments 96 offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); 97 } else { 98 CUDA_CHECK(DeviceSegmentedSort::SortPairs(nullptr, temp_storage_bytes, temp_keys, 99 temp_keys, // keys (in-place) 100 temp_indices, dst, // values (indices) 101 ncols * nrows, nrows, // num items, num segments 102 offset_iterator, offset_iterator + 1, stream)); 103 } 104 } else { 105 if (nrows == 1) { 106 CUDA_CHECK(DeviceRadixSort::SortPairsDescending(nullptr, temp_storage_bytes, temp_keys, 107 temp_keys, // keys (in-place) 108 temp_indices, dst, // values (indices) 109 ncols, 0, sizeof(float) * 8, stream)); 110 } else if (is_capturing) { 111 CUDA_CHECK(DeviceSegmentedRadixSort::SortPairsDescending( 112 nullptr, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows, 113 offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); 114 } else { 115 CUDA_CHECK(DeviceSegmentedSort::SortPairsDescending(nullptr, temp_storage_bytes, temp_keys, temp_keys, 116 temp_indices, dst, ncols * nrows, nrows, 117 offset_iterator, offset_iterator + 1, stream)); 118 } 119 } 120 121 ggml_cuda_pool_alloc<uint8_t> temp_storage_alloc(pool, temp_storage_bytes); 122 void * d_temp_storage = temp_storage_alloc.get(); 123 124 if (order == GGML_SORT_ORDER_ASC) { 125 if (nrows == 1) { 126 CUDA_CHECK(DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, 127 temp_keys, // keys (in-place) 128 temp_indices, dst, // values (indices) 129 ncols, 0, sizeof(float) * 8, stream)); 130 } else if (is_capturing) { 131 CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, 132 temp_indices, dst, ncols * nrows, nrows, offset_iterator, 133 offset_iterator + 1, 0, sizeof(float) * 8, stream)); 134 } else { 135 CUDA_CHECK(DeviceSegmentedSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, 136 temp_indices, dst, ncols * nrows, nrows, offset_iterator, 137 offset_iterator + 1, stream)); 138 } 139 } else { 140 if (nrows == 1) { 141 CUDA_CHECK(DeviceRadixSort::SortPairsDescending(d_temp_storage, temp_storage_bytes, temp_keys, 142 temp_keys, // keys (in-place) 143 temp_indices, dst, // values (indices) 144 ncols, 0, sizeof(float) * 8, stream)); 145 } else if (is_capturing) { 146 CUDA_CHECK(DeviceSegmentedRadixSort::SortPairsDescending( 147 d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows, 148 offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); 149 } else { 150 CUDA_CHECK(DeviceSegmentedSort::SortPairsDescending(d_temp_storage, temp_storage_bytes, temp_keys, 151 temp_keys, temp_indices, dst, ncols * nrows, nrows, 152 offset_iterator, offset_iterator + 1, stream)); 153 } 154 } 155} 156#endif // GGML_CUDA_USE_CUB 157 158// Bitonic sort implementation 159template<typename T> 160static inline __device__ void ggml_cuda_swap(T & a, T & b) { 161 T tmp = a; 162 a = b; 163 b = tmp; 164} 165 166template<ggml_sort_order order> 167static __global__ void k_argsort_f32_i32(const float * x, int * dst, const int ncols, int ncols_pad) { 168 // bitonic sort 169 int col = threadIdx.x; 170 int row = blockIdx.x; 171 172 if (col >= ncols_pad) { 173 return; 174 } 175 176 const float * x_row = x + row * ncols; 177 extern __shared__ int dst_row[]; 178 179 // initialize indices 180 dst_row[col] = col; 181 182 __syncthreads(); 183 184 for (int k = 2; k <= ncols_pad; k *= 2) { 185 for (int j = k / 2; j > 0; j /= 2) { 186 int ixj = col ^ j; 187 if (ixj > col) { 188 if ((col & k) == 0) { 189 if (dst_row[col] >= ncols || 190 (dst_row[ixj] < ncols && (order == GGML_SORT_ORDER_ASC ? 191 x_row[dst_row[col]] > x_row[dst_row[ixj]] : 192 x_row[dst_row[col]] < x_row[dst_row[ixj]])) 193 ) { 194 ggml_cuda_swap(dst_row[col], dst_row[ixj]); 195 } 196 } else { 197 if (dst_row[ixj] >= ncols || 198 (dst_row[col] < ncols && (order == GGML_SORT_ORDER_ASC ? 199 x_row[dst_row[col]] < x_row[dst_row[ixj]] : 200 x_row[dst_row[col]] > x_row[dst_row[ixj]])) 201 ) { 202 ggml_cuda_swap(dst_row[col], dst_row[ixj]); 203 } 204 } 205 } 206 __syncthreads(); 207 } 208 } 209 210 // copy the result to dst without the padding 211 if (col < ncols) { 212 dst[row * ncols + col] = dst_row[col]; 213 } 214} 215 216static int next_power_of_2(int x) { 217 int n = 1; 218 while (n < x) { 219 n *= 2; 220 } 221 return n; 222} 223 224void argsort_f32_i32_cuda_bitonic(const float * x, 225 int * dst, 226 const int ncols, 227 const int nrows, 228 ggml_sort_order order, 229 cudaStream_t stream) { 230 // bitonic sort requires ncols to be power of 2 231 const int ncols_pad = next_power_of_2(ncols); 232 233 const dim3 block_dims(ncols_pad, 1, 1); 234 const dim3 block_nums(nrows, 1, 1); 235 const size_t shared_mem = ncols_pad * sizeof(int); 236 237 // FIXME: this limit could be raised by ~2-4x on Ampere or newer 238 GGML_ASSERT(shared_mem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb); 239 240 if (order == GGML_SORT_ORDER_ASC) { 241 k_argsort_f32_i32<GGML_SORT_ORDER_ASC> 242 <<<block_nums, block_dims, shared_mem, stream>>>(x, dst, ncols, ncols_pad); 243 } else if (order == GGML_SORT_ORDER_DESC) { 244 k_argsort_f32_i32<GGML_SORT_ORDER_DESC> 245 <<<block_nums, block_dims, shared_mem, stream>>>(x, dst, ncols, ncols_pad); 246 } else { 247 GGML_ABORT("fatal error"); 248 } 249} 250 251void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { 252 const ggml_tensor * src0 = dst->src[0]; 253 const float * src0_d = (const float *)src0->data; 254 float * dst_d = (float *)dst->data; 255 cudaStream_t stream = ctx.stream(); 256 257 GGML_ASSERT(src0->type == GGML_TYPE_F32); 258 GGML_ASSERT( dst->type == GGML_TYPE_I32); 259 GGML_ASSERT(ggml_is_contiguous(src0)); 260 261 const int64_t ncols = src0->ne[0]; 262 const int64_t nrows = ggml_nrows(src0); 263 264 enum ggml_sort_order order = (enum ggml_sort_order) dst->op_params[0]; 265 266#ifdef GGML_CUDA_USE_CUB 267 const int ncols_pad = next_power_of_2(ncols); 268 const size_t shared_mem = ncols_pad * sizeof(int); 269 const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb; 270 271 // early return if we can use bitonic argsort 272 if (shared_mem <= max_shared_mem && ncols <= 1024) { 273 argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream); 274 return; 275 } 276 277 const int chunk_nrows = argsort_f32_i32_cuda_cub_chunk_nrows(src0->nb[1], nrows); 278 279 ggml_cuda_pool & pool = ctx.pool(); 280 281 for (int64_t i = 0; i < nrows; i += chunk_nrows) { 282 int iter_nrows = std::min((int64_t) chunk_nrows, nrows - i); 283 284 argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, iter_nrows, order, stream); 285 286 src0_d += ncols * iter_nrows; 287 dst_d += ncols * iter_nrows; 288 } 289#else 290 argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream); 291#endif 292}