char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
f8dbcd618
char/aliceai-t5-moe
1#include "argsort.cuh" 2#include "top-k.cuh" 3 4#ifdef GGML_CUDA_USE_CUB 5# include <cub/cub.cuh> 6# if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2) 7# define CUB_TOP_K_AVAILABLE 8# include <cuda/iterator> 9using namespace cub; 10# endif // CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2 11#endif // GGML_CUDA_USE_CUB 12 13#ifdef CUB_TOP_K_AVAILABLE 14 15static void top_k_cub(ggml_cuda_pool & pool, 16 const float * src, 17 int * dst, 18 const int ncols, 19 const int k, 20 cudaStream_t stream) { 21 auto requirements = cuda::execution::require(cuda::execution::determinism::not_guaranteed, 22 cuda::execution::output_ordering::unsorted); 23 auto stream_env = cuda::stream_ref{ stream }; 24 auto env = cuda::std::execution::env{ stream_env, requirements }; 25 26 auto indexes_in = cuda::make_counting_iterator(0); 27 28 size_t temp_storage_bytes = 0; 29 CUDA_CHECK(DeviceTopK::MaxPairs(nullptr, temp_storage_bytes, src, cuda::discard_iterator(), indexes_in, dst, ncols, k, 30 env)); 31 32 ggml_cuda_pool_alloc<uint8_t> temp_storage_alloc(pool, temp_storage_bytes); 33 void * d_temp_storage = temp_storage_alloc.get(); 34 35 CUDA_CHECK(DeviceTopK::MaxPairs(d_temp_storage, temp_storage_bytes, src, cuda::discard_iterator(), indexes_in, dst, 36 ncols, k, env)); 37} 38 39#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE 40 41static int next_power_of_2(int x) { 42 int n = 1; 43 while (n < x) { 44 n *= 2; 45 } 46 return n; 47} 48 49#endif // CUB_TOP_K_AVAILABLE 50 51#if !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP) 52 53static __device__ __forceinline__ uint32_t top_k_float_to_ordered(float value) { 54 const uint32_t bits = __float_as_uint(value); 55 const uint32_t mask = (uint32_t) (-(int32_t) (bits >> 31)) | 0x80000000U; 56 return bits ^ mask; 57} 58 59struct top_k_radix_state { 60 uint32_t prefix; 61 uint32_t prefix_mask; 62 int rank; 63 int greater_count; 64 int equal_count; 65}; 66 67static __global__ void top_k_radix_init(top_k_radix_state * states, int nrows, int k) { 68 const int row = blockIdx.x * blockDim.x + threadIdx.x; 69 if (row < nrows) { 70 states[row] = {0, 0, k, 0, 0}; 71 } 72} 73 74template<int BLOCK_SIZE, int RADIX_BITS> 75static __global__ void top_k_radix_histogram( 76 const float * __restrict__ src, 77 const top_k_radix_state * __restrict__ states, 78 int * __restrict__ block_histograms, 79 int ncols, 80 int blocks_per_row, 81 int shift) { 82 constexpr int NBINS = 1 << RADIX_BITS; 83 84 const int row = blockIdx.x / blocks_per_row; 85 const int row_block = blockIdx.x % blocks_per_row; 86 const int tid = threadIdx.x; 87 const float * row_src = src + (size_t) row * ncols; 88 __shared__ int histogram[NBINS]; 89 90 histogram[tid] = 0; 91 __syncthreads(); 92 93 const top_k_radix_state state = states[row]; 94 for (int col = row_block * BLOCK_SIZE + tid; 95 col < ncols; 96 col += blocks_per_row * BLOCK_SIZE) { 97 const uint32_t key = top_k_float_to_ordered(row_src[col]); 98 if ((key & state.prefix_mask) == state.prefix) { 99 atomicAdd(&histogram[(key >> shift) & (NBINS - 1)], 1); 100 } 101 } 102 __syncthreads(); 103 104 const size_t histogram_offset = 105 ((size_t) row * blocks_per_row + row_block) * NBINS; 106 block_histograms[histogram_offset + tid] = histogram[tid]; 107} 108 109template<int BLOCK_SIZE, int RADIX_BITS> 110static __global__ void top_k_radix_select( 111 const int * __restrict__ block_histograms, 112 top_k_radix_state * __restrict__ states, 113 int blocks_per_row, 114 int shift) { 115 constexpr int NBINS = 1 << RADIX_BITS; 116 117 const int row = blockIdx.x; 118 const int tid = threadIdx.x; 119 __shared__ int histogram[NBINS]; 120 121 int count = 0; 122 for (int row_block = 0; row_block < blocks_per_row; ++row_block) { 123 const size_t offset = ((size_t) row * blocks_per_row + row_block) * NBINS; 124 count += block_histograms[offset + tid]; 125 } 126 histogram[tid] = count; 127 __syncthreads(); 128 129 if (tid == 0) { 130 top_k_radix_state state = states[row]; 131 int bin = NBINS - 1; 132 while (bin > 0 && histogram[bin] < state.rank) { 133 state.rank -= histogram[bin--]; 134 } 135 state.prefix |= (uint32_t) bin << shift; 136 state.prefix_mask |= (uint32_t) (NBINS - 1) << shift; 137 states[row] = state; 138 } 139} 140 141static __global__ void top_k_radix_reset_counters(top_k_radix_state * states, int nrows) { 142 const int row = blockIdx.x * blockDim.x + threadIdx.x; 143 if (row < nrows) { 144 states[row].greater_count = 0; 145 states[row].equal_count = 0; 146 } 147} 148 149template<int BLOCK_SIZE> 150static __global__ void top_k_radix_gather( 151 const float * __restrict__ src, 152 int * __restrict__ dst, 153 top_k_radix_state * __restrict__ states, 154 int ncols, 155 int k, 156 int blocks_per_row) { 157 const int row = blockIdx.x / blocks_per_row; 158 const int row_block = blockIdx.x % blocks_per_row; 159 const int tid = threadIdx.x; 160 const float * row_src = src + (size_t) row * ncols; 161 int * row_dst = dst + (size_t) row * k; 162 top_k_radix_state * state = &states[row]; 163 164 for (int col = row_block * BLOCK_SIZE + tid; 165 col < ncols; 166 col += blocks_per_row * BLOCK_SIZE) { 167 const uint32_t key = top_k_float_to_ordered(row_src[col]); 168 if (key > state->prefix) { 169 const int pos = atomicAdd(&state->greater_count, 1); 170 row_dst[pos] = col; 171 } else if (key == state->prefix) { 172 const int pos = atomicAdd(&state->equal_count, 1); 173 if (pos < state->rank) { 174 row_dst[k - state->rank + pos] = col; 175 } 176 } 177 } 178} 179 180static void top_k_radix_cuda( 181 ggml_cuda_pool & pool, 182 const float * src, int * dst, int ncols, int nrows, int k, cudaStream_t stream) { 183 constexpr int BLOCK_SIZE = 256; 184 constexpr int RADIX_BITS = 8; 185 constexpr int NBINS = 1 << RADIX_BITS; 186 const int blocks_per_row = std::min((ncols + 1023) / 1024, 64); 187 188 ggml_cuda_pool_alloc<top_k_radix_state> states_alloc(pool, nrows); 189 ggml_cuda_pool_alloc<int> histograms_alloc(pool, (size_t) nrows * blocks_per_row * NBINS); 190 top_k_radix_state * states = states_alloc.get(); 191 int * histograms = histograms_alloc.get(); 192 193 top_k_radix_init<<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows, k); 194 195 const dim3 row_grid(blocks_per_row * nrows); 196 for (int shift = 32 - RADIX_BITS; shift >= 0; shift -= RADIX_BITS) { 197 top_k_radix_histogram<BLOCK_SIZE, RADIX_BITS> 198 <<<row_grid, BLOCK_SIZE, 0, stream>>>( 199 src, states, histograms, ncols, blocks_per_row, shift); 200 top_k_radix_select<BLOCK_SIZE, RADIX_BITS> 201 <<<nrows, BLOCK_SIZE, 0, stream>>>(histograms, states, blocks_per_row, shift); 202 } 203 204 top_k_radix_reset_counters 205 <<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows); 206 top_k_radix_gather<BLOCK_SIZE> 207 <<<row_grid, BLOCK_SIZE, 0, stream>>>( 208 src, dst, states, ncols, k, blocks_per_row); 209} 210 211#endif // !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP) 212 213void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { 214 const ggml_tensor * src0 = dst->src[0]; 215 const float * src0_d = (const float *) src0->data; 216 int * dst_d = (int *) dst->data; 217 cudaStream_t stream = ctx.stream(); 218 219 // are these asserts truly necessary? 220 GGML_ASSERT(src0->type == GGML_TYPE_F32); 221 GGML_ASSERT(dst->type == GGML_TYPE_I32); 222 GGML_ASSERT(ggml_is_contiguous(src0)); 223 224 const int64_t ncols = src0->ne[0]; 225 const int64_t nrows = ggml_nrows(src0); 226 const int64_t k = dst->ne[0]; 227 ggml_cuda_pool & pool = ctx.pool(); 228#ifdef CUB_TOP_K_AVAILABLE 229 // TODO: Switch to `DeviceSegmentedTopK` for multi-row TopK once implemented 230 // https://github.com/NVIDIA/cccl/issues/6391 231 // TODO: investigate if there exists a point where parallelized argsort is faster than sequential top-k 232 for (int i = 0; i < nrows; i++) { 233 top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream); 234 } 235#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE 236 // Fall back to argsort + copy 237 const int ncols_pad = next_power_of_2(ncols); 238 const size_t shared_mem = ncols_pad * sizeof(int); 239 const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb; 240 const bool use_bitonic = shared_mem <= max_shared_mem && ncols <= 1024; 241 const int chunk_nrows = argsort_f32_i32_cuda_cub_chunk_nrows(src0->nb[1], nrows); 242 243 ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * chunk_nrows); 244 int * tmp_dst = temp_dst_alloc.get(); 245 246 for (int64_t i = 0; i < nrows; i += chunk_nrows) { 247 int iter_nrows = std::min((int64_t) chunk_nrows, nrows - i); 248 249 if (use_bitonic) { 250 argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, iter_nrows, GGML_SORT_ORDER_DESC, stream); 251 } else { 252 argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, iter_nrows, GGML_SORT_ORDER_DESC, stream); 253 } 254 CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), iter_nrows, 255 cudaMemcpyDeviceToDevice, stream)); 256 257 src0_d += ncols * iter_nrows; 258 dst_d += k * iter_nrows; 259 } 260#else // GGML_CUDA_USE_CUB 261#if defined(GGML_USE_HIP) 262 if (ncols > 1024) { 263 top_k_radix_cuda(pool, src0_d, dst_d, ncols, nrows, k, stream); 264 } else { 265#endif // defined(GGML_USE_HIP) 266 ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows); 267 int * tmp_dst = temp_dst_alloc.get(); 268 argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream); 269 CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows, 270 cudaMemcpyDeviceToDevice, stream)); 271#if defined(GGML_USE_HIP) 272 } 273#endif // defined(GGML_USE_HIP) 274#endif 275}