diff --git a/ggml/src/ggml-cuda/lightning-indexer.cu b/ggml/src/ggml-cuda/lightning-indexer.cu index eee10ca215..6060972361 100644 --- a/ggml/src/ggml-cuda/lightning-indexer.cu +++ b/ggml/src/ggml-cuda/lightning-indexer.cu @@ -4,238 +4,6 @@ #include "convert.cuh" #include "mma.cuh" -#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) -#if defined(TURING_MMA_AVAILABLE) - -typedef union { - int2 i2; - half2 h2[2]; -} half4; - -#include -namespace wmma = nvcuda::wmma; - -template -static __global__ void lightning_indexer_kernel_wmma( - 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 HEADS_PER_INNER_LOOP = 8; - constexpr int K_EMBD_PER_INNER_LOOP = 16; - constexpr int N_EMBD_PADDED = N_EMBD + 8; - - const int i_batch = blockIdx.y; - const int i_stream = blockIdx.z; - const int i_warp = threadIdx.y; - const int i_lane = threadIdx.x; - const int tid = i_warp * WARP_SIZE + i_lane; - - // each block processes K_VECS_PER_BLOCK K vectors - const int start_kv = blockIdx.x * K_VECS_PER_BLOCK; - - const char * q_base = (const char *) Q + i_batch*nbq2 + i_stream*nbq3; - const float * w_base = (const float *) ((const char *) W + i_batch*nbw1 + i_stream*nbw3); - - // phase 1 - load weights and first Q tile to shared memory - - __shared__ float w_shared[N_HEAD]; - __shared__ int2 q_shared_h[HEADS_PER_INNER_LOOP][N_EMBD_PADDED / 4]; - - if (tid < N_HEAD) { - w_shared[tid] = w_base[tid]; - } - - // total number of half4 elements in HEADS_PER_INNER_LOOP x N_EMBD Q tile - constexpr int N_Q_TILE = HEADS_PER_INNER_LOOP * (N_EMBD / 4); - // number of registers needed in each thread to store Q tile in thread block - constexpr int N_Q_NEXT = (N_Q_TILE + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; - -#pragma unroll - for (int i_q = tid; i_q < N_Q_TILE; i_q += THREADS_PER_BLOCK) { - const int i_head = i_q / (N_EMBD / 4); - const int i_embd = i_q % (N_EMBD / 4); - const float4 q = *(const float4 *) (q_base + i_head*nbq1 + i_embd*sizeof(float4)); - half4 q_packed; - q_packed.h2[0] = __float22half2_rn(make_float2(q.x, q.y)); - q_packed.h2[1] = __float22half2_rn(make_float2(q.z, q.w)); - q_shared_h[i_head][i_embd] = q_packed.i2; - } - - // phase 2 - load (and dequantize if needed) K to shared mem - - __shared__ half2 k_shared_h[K_VECS_PER_BLOCK][N_EMBD_PADDED / 4][2]; - - constexpr int n_k = K_VECS_PER_BLOCK * (N_EMBD / 4); - - if constexpr (TYPE_K == GGML_TYPE_F16) { -#pragma unroll - for (int i_k = tid; i_k < n_k; i_k += THREADS_PER_BLOCK) { - const int i_k_vec = i_k / (N_EMBD / 4); - const int i_embd = i_k % (N_EMBD / 4); - const int i_kv = start_kv + i_k_vec; - if (i_kv < n_kv) { - const int2 * k_base = (const int2 *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3); - *(int2*) &k_shared_h[i_k_vec][i_embd] = k_base[i_embd]; - } else { - *(int2*) &k_shared_h[i_k_vec][i_embd] = make_int2(0, 0); - } - } - } else { - constexpr dequantize_V_t dequantize_k = get_dequantize_V(); -#pragma unroll - for (int i_k = tid; i_k < n_k; i_k += THREADS_PER_BLOCK) { - const int i_k_vec = i_k / (N_EMBD / 4); - const int i_embd = i_k % (N_EMBD / 4); - const int i_kv = start_kv + i_k_vec; - if (i_kv < n_kv) { - const void * k_base = (const void *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3); - dequantize_k(k_base, &k_shared_h[i_k_vec][i_embd][0], i_embd * 4); - } else { - *(int2*) &k_shared_h[i_k_vec][i_embd] = make_int2(0, 0); - } - } - } - - __syncthreads(); - - // phase 3 - calculate lightning indexer scores - - __shared__ float qk_shared[WARPS_PER_BLOCK][HEADS_PER_INNER_LOOP][K_VECS_PER_BLOCK]; - - // load K fragment - wmma::fragment frag_k; - wmma::load_matrix_sync(frag_k, (half*) &k_shared_h[0][i_warp * K_EMBD_PER_INNER_LOOP / 4], N_EMBD_PADDED); - - float score_k = 0.0f; - - for (int i_head_0 = 0; i_head_0 < N_HEAD; i_head_0 += HEADS_PER_INNER_LOOP) { - const int i_head_next = i_head_0 + HEADS_PER_INNER_LOOP; - - // we don't use accumulator for anything, fill it with zeros - wmma::fragment frag_acc; - wmma::fill_fragment(frag_acc, 0.0f); - - // load Q fragment - wmma::fragment frag_q; - wmma::load_matrix_sync(frag_q, (half*) &q_shared_h[0][i_warp * K_EMBD_PER_INNER_LOOP / 4], N_EMBD_PADDED); - - // preload next Q tile to registers during matrix multiplication - float4 q_next[N_Q_NEXT]; - - if (i_head_next < N_HEAD) { -#pragma unroll - for (int i_q = tid, i_q_next = 0; i_q < N_Q_TILE; i_q += THREADS_PER_BLOCK) { - const int i_head = i_head_next + i_q / (N_EMBD / 4); - const int i_embd = i_q % (N_EMBD / 4); - q_next[i_q_next++] = *(const float4 *) (q_base + i_head*nbq1 + i_embd*sizeof(float4)); - } - } - - // perform matrix multiplication - wmma::mma_sync(frag_acc, frag_q, frag_k, frag_acc); - wmma::store_matrix_sync((float*) &qk_shared[i_warp][0][0], frag_acc, K_VECS_PER_BLOCK, wmma::mem_row_major); - - // make sure all threads finished using q_shared_h so we can store next tile - __syncthreads(); - - // write preloaded Q tile to shared memory - if (i_head_next < N_HEAD) { -#pragma unroll - for (int i_q = tid, i_q_next = 0; i_q < N_Q_TILE; i_q += THREADS_PER_BLOCK) { - const int i_head = i_q / (N_EMBD / 4); - const int i_embd = i_q % (N_EMBD / 4); - half4 q_packed; - q_packed.h2[0] = __float22half2_rn(make_float2(q_next[i_q_next].x, q_next[i_q_next].y)); - q_packed.h2[1] = __float22half2_rn(make_float2(q_next[i_q_next].z, q_next[i_q_next].w)); - q_shared_h[i_head][i_embd] = q_packed.i2; - ++i_q_next; - } - } - - // accumulate QK multiplication results from all block warps - // (there are 256 threads in block and 256 matmul outputs) - // TODO it will break if WARP_SIZE is not 32 - const int h = tid / K_VECS_PER_BLOCK; - const int k = tid % K_VECS_PER_BLOCK; - const float w_val = w_shared[i_head_0 + h]; - - float sum = 0.0f; -#pragma unroll - for (int w = 0; w < WARPS_PER_BLOCK; ++w) { - sum += qk_shared[w][h][k]; - } - - // ReLU, weight - sum = sum > 0.0f ? sum : 0.0f; - sum *= w_val; - - // wait until qk_shared[0] is no longer used - __syncthreads(); - - // reuse qk_shared[0] for storing partial results - qk_shared[0][h][k] = sum; - - // wait until all threads write their results - __syncthreads(); - - // accumulate result over heads - if (tid < K_VECS_PER_BLOCK) { -#pragma unroll - for (int i_head = 0; i_head < HEADS_PER_INNER_LOOP; ++i_head) { - score_k += qk_shared[0][i_head][tid]; - } - } - - // make sure all threads finished using qk_shared - __syncthreads(); - } - - // phase 4 - store output to VRAM - - if (tid < K_VECS_PER_BLOCK) { - const int i_kv = start_kv + tid; - if (i_kv < n_kv) { - 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); - dst_base[i_kv] = score_k + __half2float(m_base[i_kv]); - } - } -} - -#else // defined(TURING_MMA_AVAILABLE) - -template -static __global__ void lightning_indexer_kernel_wmma( - 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 - ) { - GGML_UNUSED_VARS(Q, K, W, M, dst, - n_stream, n_batch, n_kv, - nb1, nb2, nb3, - nbq1, nbq2, nbq3, - nbk1, nbk2, nbk3, - nbw1, nbw2, nbw3, - nem3); - NO_DEVICE_CODE; -} - -#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 @@ -816,27 +584,22 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * GGML_ABORT("fatal error"); } } else if (n_embd == 128 && n_head == 32) { -#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) - if (GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16) { - // use wmma kernel + if ((turing_mma_available(cc) || amd_mfma_available(cc) || amd_wmma_available(cc)) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16) { constexpr int K_VECS_PER_BLOCK = 32; - constexpr int WARPS_PER_BLOCK = 8; + constexpr int WARPS_PER_BLOCK = 4; - dim3 block(32, WARPS_PER_BLOCK); + dim3 block(ggml_cuda_info().devices[device].warp_size, WARPS_PER_BLOCK); int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK); dim3 grid(num_kv_blocks, n_batch, n_stream); - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_F16) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q4_0) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q4_1) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q5_0) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q5_1) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q8_0) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 32, k, GGML_TYPE_F16) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 32, k, GGML_TYPE_Q4_0) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 32, k, GGML_TYPE_Q4_1) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 32, k, GGML_TYPE_Q5_0) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 32, k, GGML_TYPE_Q5_1) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 32, k, GGML_TYPE_Q8_0) GGML_ABORT("fatal error"); } else { -#else // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) - { -#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) // use vector kernel constexpr int K_VECS_PER_WARP = 8; constexpr int WARPS_PER_BLOCK = 8;