public_git/whisper.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/whisper.cpp.git synced 2026-10-20 17:23:00 +02:00
Code Issues Packages Projects Releases Wiki Activity

musa: use the tile lightning indexer kernel (llama/30080)

Browse Source 
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
R0CKSTAR
2026-10-19 09:47:10 +02:00
committed by Daniel Bevenius
parent 7f56abab92
commit 298d79dd9e
1 changed files with 53 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/lightning-indexer.cu
+53 -36
View File
@@ -239,6 +239,14 @@ static __global__ void lightning_indexer_kernel_wmma(
// tokens scored per block by the tile kernel
#define LIGHTNING_INDEXER_TILE_TOKENS 8
// heads whose queries the tile kernel stages per pass, MUSA arch 21 caps static shared memory
// at 28 KB and the queries of four heads do not fit there next to the key tile
#if defined(GGML_USE_MUSA) && defined(__MUSA_ARCH__) && __MUSA_ARCH__ < 220
#define LIGHTNING_INDEXER_TILE_HEADS_PER_PASS 2
#else
#define LIGHTNING_INDEXER_TILE_HEADS_PER_PASS 4
#endif
// 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)
@@ -406,9 +414,11 @@ static __global__ void lightning_indexer_kernel_tile(
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;
constexpr int HEADS_PER_PASS = N_HEAD < LIGHTNING_INDEXER_TILE_HEADS_PER_PASS ? N_HEAD : LIGHTNING_INDEXER_TILE_HEADS_PER_PASS;
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");
static_assert(N_HEAD % HEADS_PER_PASS == 0, "head passes must cover the heads");
const int tid = threadIdx.y * WARP_SIZE + threadIdx.x;
const int start_kv = blockIdx.x * K_VECS_PER_BLOCK;
@@ -417,7 +427,7 @@ static __global__ void lightning_indexer_kernel_tile(
// 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__ float2 q_shared[HEADS_PER_PASS][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
@@ -451,22 +461,7 @@ static __global__ void lightning_indexer_kernel_tile(
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);
}
// phase 2 - stage the weights of every head, tokens past n_batch are zero
if (tid < N_HEAD * TOKENS_PER_BLOCK) {
const int h = tid / TOKENS_PER_BLOCK;
@@ -475,33 +470,60 @@ static __global__ void lightning_indexer_kernel_tile(
((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]);
for (int h0 = 0; h0 < N_HEAD; h0 += HEADS_PER_PASS) {
// the previous pass is fully consumed before its queries are replaced
if (h0 > 0) {
__syncthreads();
}
// phase 3 - stage the queries of the heads of this pass, tokens past n_batch are zero
#pragma unroll
for (int h = 0; h < N_HEAD; ++h) {
const float2 q_val = q_shared[h][tl][c];
for (int i = tid; i < HEADS_PER_PASS * 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 + (h0 + 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);
}
__syncthreads();
// phase 4 - float products of the widened keys for the heads of this pass
#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) {
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]);
k_val[j] = __half22float2(k_shared[kl + j*KEY_LANES][c]);
}
#pragma unroll
for (int h = 0; h < HEADS_PER_PASS; ++h) {
const float2 q_val = q_shared[h][tl][c];
#pragma unroll
for (int j = 0; j < KEYS_PER_THREAD; ++j) {
qk[h0 + h][j] = fmaf(k_val[j].x, q_val.x, qk[h0 + h][j]);
qk[h0 + h][j] = fmaf(k_val[j].y, q_val.y, qk[h0 + h][j]);
}
}
}
}
// phase 5 - ReLU, weight, add the mask and write, consecutive threads write consecutive keys
float score[KEYS_PER_THREAD] = { 0.0f };
#pragma unroll
@@ -512,8 +534,6 @@ static __global__ void lightning_indexer_kernel_tile(
}
}
// 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;
@@ -677,8 +697,6 @@ 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");
}
#ifndef GGML_USE_MUSA
// MUSA archs 21 and 22 cap static shared memory at 28 KB, below what the tile kernel stages
} 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;
@@ -698,9 +716,8 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor *
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");
#endif // GGML_USE_MUSA
} else if (n_embd == 128 && n_head == 4) {
// a batch smaller than a token tile, or MUSA, 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;