mirrors/llama.cpp
0
0 Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-10-28 17:56:13 +08:00
Code Activity

CUDA: Remove wmma lightning-indexer, also run head 32 through mma kernel

Browse Source
This commit is contained in:
Carl Philipp Klemm committed 2026-10-25 11:39:29 +01:00
1 parent bcd5ac991f
commit 90ed669d35
1 file changed
+9 -246
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/lightning-indexer.cu
+9 -246
View File
@@ -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 <mma.h>
namespace wmma = nvcuda::wmma;
template <int WARPS_PER_BLOCK, int K_VECS_PER_BLOCK, int64_t N_EMBD, int64_t N_HEAD, ggml_type TYPE_K>
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<TYPE_K, half, 4>();
#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<wmma::matrix_b, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, half, wmma::col_major> 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<wmma::accumulator, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, float> frag_acc;
wmma::fill_fragment(frag_acc, 0.0f);
// load Q fragment
wmma::fragment<wmma::matrix_a, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, half, wmma::row_major> 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 <int WARPS_PER_BLOCK, int K_VECS_PER_BLOCK, int64_t N_EMBD, int64_t N_HEAD, ggml_type TYPE_K>
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;