diff --git a/ggml/src/ggml-cuda/lightning-indexer.cu b/ggml/src/ggml-cuda/lightning-indexer.cu index 54e0e6314..90c06a1cf 100644 --- a/ggml/src/ggml-cuda/lightning-indexer.cu +++ b/ggml/src/ggml-cuda/lightning-indexer.cu @@ -236,6 +236,9 @@ static __global__ void lightning_indexer_kernel_wmma( #endif // defined(TURING_MMA_AVAILABLE) #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +// tokens scored per block by the tile kernel +#define LIGHTNING_INDEXER_TILE_TOKENS 8 + // TODO there is one ugly assumption used in this kernel - that WARP_SIZE is equal to 32 // thanks to that one warp operating on float4 processes whole indexer K/Q vectors // 32 * 4 = 128 (N_EMBD) @@ -382,6 +385,152 @@ static __global__ void lightning_indexer_kernel_vec( } } +// one block scores a tile of K_VECS_PER_BLOCK keys against TOKENS_PER_BLOCK tokens: the keys are +// staged in half precision and the queries of every head in float, each thread owns KEYS_PER_THREAD +// keys for one token, a warp shares its token so the query reads are broadcasts, and every key +// element is widened once for all heads, so no dot product needs a cross thread reduction +template +static __global__ void lightning_indexer_kernel_tile( + const float * Q, const char * K, const float * W, const half * M, float * dst, + int64_t n_stream, int64_t n_batch, int64_t n_kv, + size_t nb1, size_t nb2, size_t nb3, + size_t nbq1, size_t nbq2, size_t nbq3, + size_t nbk1, size_t nbk2, size_t nbk3, + size_t nbw1, size_t nbw2, size_t nbw3, + size_t nbm1, size_t nbm2, size_t nbm3, + int64_t nem3 + ) { + + constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WARP_SIZE; + constexpr int TOKENS_PER_BLOCK = LIGHTNING_INDEXER_TILE_TOKENS; + constexpr int KEY_LANES = THREADS_PER_BLOCK / TOKENS_PER_BLOCK; + constexpr int KEYS_PER_THREAD = K_VECS_PER_BLOCK / KEY_LANES; + constexpr int N_EMBD_H2 = N_EMBD / 2; + + static_assert(THREADS_PER_BLOCK % TOKENS_PER_BLOCK == 0, "threads must cover the token tile"); + static_assert(K_VECS_PER_BLOCK % KEY_LANES == 0, "key lanes must cover the key tile"); + + const int tid = threadIdx.y * WARP_SIZE + threadIdx.x; + const int start_kv = blockIdx.x * K_VECS_PER_BLOCK; + const int start_batch = blockIdx.y * TOKENS_PER_BLOCK; + const int i_stream = blockIdx.z; + + // the row padding keeps the keys of consecutive threads in distinct banks + __shared__ half2 k_shared[K_VECS_PER_BLOCK][N_EMBD_H2 + 1]; + __shared__ float2 q_shared[N_HEAD][TOKENS_PER_BLOCK][N_EMBD_H2]; + __shared__ float w_shared[N_HEAD][TOKENS_PER_BLOCK]; + + // phase 1 - stage the key tile four elements at a time, rows past n_kv are zero + +#pragma unroll + for (int i = tid; i < K_VECS_PER_BLOCK * (N_EMBD / 4); i += THREADS_PER_BLOCK) { + const int r = i / (N_EMBD / 4); + const int c4 = i % (N_EMBD / 4); + + half2 lo = make_half2(0.0f, 0.0f); + half2 hi = lo; + if (start_kv + r < n_kv) { + const char * k_row = K + (start_kv + r)*nbk2 + i_stream*nbk3; + if constexpr (TYPE_K == GGML_TYPE_F16) { + lo = ((const half2 *) k_row)[2*c4 + 0]; + hi = ((const half2 *) k_row)[2*c4 + 1]; + } else { + float4 v; + if constexpr (TYPE_K == GGML_TYPE_F32) { + v = ((const float4 *) k_row)[c4]; + } else { + constexpr dequantize_V_t dequantize_k = get_dequantize_V(); + dequantize_k(k_row, &v, c4 * 4); + } + lo = make_half2(v.x, v.y); + hi = make_half2(v.z, v.w); + } + } + + k_shared[r][2*c4 + 0] = lo; + k_shared[r][2*c4 + 1] = hi; + } + + // phase 2 - stage the queries and weights of every head, tokens past n_batch are zero + +#pragma unroll + for (int i = tid; i < N_HEAD * TOKENS_PER_BLOCK * (N_EMBD / 4); i += THREADS_PER_BLOCK) { + const int h = i / (TOKENS_PER_BLOCK * (N_EMBD / 4)); + const int r = i / (N_EMBD / 4) % TOKENS_PER_BLOCK; + const int c4 = i % (N_EMBD / 4); + + float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + if (start_batch + r < n_batch) { + v = *(const float4 *) ((const char *) Q + h*nbq1 + (start_batch + r)*nbq2 + i_stream*nbq3 + c4*sizeof(float4)); + } + + q_shared[h][r][2*c4 + 0] = make_float2(v.x, v.y); + q_shared[h][r][2*c4 + 1] = make_float2(v.z, v.w); + } + + if (tid < N_HEAD * TOKENS_PER_BLOCK) { + const int h = tid / TOKENS_PER_BLOCK; + const int r = tid % TOKENS_PER_BLOCK; + w_shared[h][r] = start_batch + r < n_batch ? + ((const float *) ((const char *) W + (start_batch + r)*nbw1 + i_stream*nbw3))[h] : 0.0f; + } + + __syncthreads(); + + // phase 3 - float products of the widened keys for every head, ReLU, weight + + const int kl = tid % KEY_LANES; + const int tl = tid / KEY_LANES; + + float qk[N_HEAD][KEYS_PER_THREAD] = { { 0.0f } }; + +#pragma unroll 8 + for (int c = 0; c < N_EMBD_H2; ++c) { + float2 k_val[KEYS_PER_THREAD]; +#pragma unroll + for (int j = 0; j < KEYS_PER_THREAD; ++j) { + k_val[j] = __half22float2(k_shared[kl + j*KEY_LANES][c]); + } +#pragma unroll + for (int h = 0; h < N_HEAD; ++h) { + const float2 q_val = q_shared[h][tl][c]; +#pragma unroll + for (int j = 0; j < KEYS_PER_THREAD; ++j) { + qk[h][j] = fmaf(k_val[j].x, q_val.x, qk[h][j]); + qk[h][j] = fmaf(k_val[j].y, q_val.y, qk[h][j]); + } + } + } + + float score[KEYS_PER_THREAD] = { 0.0f }; + +#pragma unroll + for (int h = 0; h < N_HEAD; ++h) { +#pragma unroll + for (int j = 0; j < KEYS_PER_THREAD; ++j) { + score[j] += fmaxf(qk[h][j], 0.0f) * w_shared[h][tl]; + } + } + + // phase 4 - add the mask and write, consecutive threads write consecutive keys + + const int i_batch = start_batch + tl; + if (i_batch >= n_batch) { + return; + } + + const half * m_base = (const half *) ((const char *) M + i_batch*nbm1 + (i_stream%nem3)*nbm3); + float * dst_base = (float *) ((char *) dst + i_batch*nb1 + i_stream*nb3); + +#pragma unroll + for (int j = 0; j < KEYS_PER_THREAD; ++j) { + const int i_kv = start_kv + kl + j*KEY_LANES; + if (i_kv < n_kv) { + dst_base[i_kv] = score[j] + __half2float(m_base[i_kv]); + } + } +} + #define LIGHTNING_INDEXER_CASE(lightning_indexer_kernel, n_embd, n_head, K, type_K) \ if (K->type == (type_K)) { \ lightning_indexer_kernel \ @@ -528,8 +677,27 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_F32) GGML_ABORT("fatal error"); } + } else if (n_embd == 128 && n_head == 4 && n_batch >= LIGHTNING_INDEXER_TILE_TOKENS) { + // too few heads for a wmma tile, the tile kernel shares the keys across the tokens + constexpr int WARPS_PER_BLOCK = 8; + constexpr int K_VECS_PER_BLOCK = 64; + + dim3 block(32, WARPS_PER_BLOCK); + int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK); + int num_batch_blocks = (n_batch + LIGHTNING_INDEXER_TILE_TOKENS - 1) / LIGHTNING_INDEXER_TILE_TOKENS; + dim3 grid(num_kv_blocks, num_batch_blocks, n_stream); + + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_F16) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_Q4_0) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_Q4_1) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_Q5_0) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_Q5_1) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_Q8_0) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_BF16) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_F32) + GGML_ABORT("fatal error"); } else if (n_embd == 128 && n_head == 4) { - // too few heads for a wmma tile, use vector kernel + // a batch smaller than a token tile, use vector kernel constexpr int K_VECS_PER_WARP = 8; constexpr int WARPS_PER_BLOCK = 8; constexpr int K_VECS_PER_BLOCK = K_VECS_PER_WARP * WARPS_PER_BLOCK;