public_git
/
whisper.cpp
mirror of https://github.com/ggml-org/whisper.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

opencl: transpose the K tile in local memory for FA prefill kernels (llama/26428)

This commit is contained in:
Hongqiang Wang 2026-08-10 11:09:19 -07:00 committed by Georgi Gerganov
parent 830ec22572
commit a6e2630ff1
4 changed files with 156 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
ggml/src/ggml-opencl/ggml-opencl.cpp
View File

@ -73,6 +73,7 @@ typedef const void * (*get_adreno_bin_kernel_func_t)(
//------------------------------------------------------------------------------
bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor);
static bool ggml_cl_is_q4_0_soa(const ggml_tensor * tensor);
static bool ggml_cl_is_q8_0_soa(const ggml_tensor * tensor);
static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);
@ -4629,6 +4630,23 @@ static std::string ggml_opencl_fa_compile_opts(ggml_backend_opencl_context * bac
if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) {
opts += " -D FA_C8_NO_SG_PIN";
}
// Transposed K tile in local memory: the KV rows the QK loop walks together become
// adjacent, so a group of them is ONE 128-bit local read instead of several narrow
// ones. The QK loop is LDS-read-issue-bound (a wrong-math probe that kept every FMA/dp4a
// but removed the LDS reads ran the kernel ~40% faster), so this is worth up to +26% on
// fa=1 prefill. Output is bit-identical -- only the layout moves.
//
// DK <= 128 only. At DK=256 (gemma-3-4b) it measures 1-2% NEGATIVE and reproduces across
// rounds; padding the row stride does not recover it, so the cause is not a simple bank
// conflict and the wider tile does not want this layout.
//
// Default on within that gate; GGML_OPENCL_FA_K_LDS_T=0 restores the row-major tile.
{
const char * e = getenv("GGML_OPENCL_FA_K_LDS_T");
if ((e == nullptr || e[0] != '0') && cfg->dk <= 128) {
opts += " -D FA_K_LDS_T";
}
}
return opts;
}

51
ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl
View File

@ -211,7 +211,30 @@ __kernel void FA_TILE_NAME(
float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
#ifdef FA_K_LDS_T
// K tile transposed: [dk vec][kv row] instead of [kv row][dk vec].
//
// The QK loop walks 2 or 4 KV rows at a time against the same dk element. Row-major
// those are DK_VEC half4s apart, so each is its own 64-bit local read. Transposed they
// are adjacent, so a pair is one 128-bit read -- half the LDS issues for the same bytes,
// no extra registers, arithmetic untouched.
//
// This kernel looked like it should be FMA-bound (a half4 mad does ~4 ALU ops per LDS
// read, unlike the 1:1 of the dp4a loop), but it is NOT: a wrong-math probe that kept
// every FMA and removed the LDS reads ran it 38.6% faster (18.92 -> 11.62 ms/op).
// Explicitly 16-byte aligned: FA_LK_PAIR below reads two adjacent half4 as one float4,
// and the element type only obliges the compiler to align this array to 8. The indices
// are even so the offset is a multiple of 16, but the base has to be too, and relying
// on the compiler to over-align it is relying on luck.
__local KV_DATA_TYPE4 l_k[DK_VEC][BLOCK_N] __attribute__((aligned(16)));
#define FA_LK(ROW, C) l_k[C][ROW]
// Two adjacent KV rows as one 128-bit local read (half4 pair == 16 B). j is even and
// BLOCK_N is even, so &l_k[c][j] is 16 B past a 16 B-aligned base.
#define FA_LK_PAIR(C, J) as_half8(*(__local const float4 *)(&l_k[C][J]))
#else
__local KV_DATA_TYPE4 l_k[BLOCK_N][DK_VEC];
#define FA_LK(ROW, C) l_k[ROW][C]
#endif
__local KV_DATA_TYPE4 l_v[BLOCK_N][DV_VEC];
#if N_SPLIT > 1 && !defined(HAS_SUBGROUP_SHUFFLE)
@ -254,17 +277,17 @@ __kernel void FA_TILE_NAME(
#ifdef FA_K_IMG
if (use_kv_pad) {
const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1;
l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
FA_LK(row, col) = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
} else {
const int k_row_px = batch_idx * k_pitch_px_batch + head_kv_idx * k_pitch_px_head + k_row_idx * k_pitch_px_row;
l_k[row][col] = read_imageh(k_img, k_row_px + col);
FA_LK(row, col) = read_imageh(k_img, k_row_px + col);
}
#else
const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1;
l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
FA_LK(row, col) = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
#endif
} else {
l_k[row][col] = (KV_DATA_TYPE4)(0.0h);
FA_LK(row, col) = (KV_DATA_TYPE4)(0.0h);
}
}
for (int i = tid; i < BLOCK_N * DV_VEC; i += WG_SIZE) {
@ -292,8 +315,15 @@ __kernel void FA_TILE_NAME(
FA_UNROLL
for (int k = 0; k < SPLIT_DK_VEC; k++) {
const ACC_TYPE4 qk = q_priv[k];
#if defined(FA_K_LDS_T)
// 2 KV rows adjacent in the transposed tile: one 128-bit local read.
const half8 kk = FA_LK_PAIR(dk_off + k, j);
ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(kk.lo);
ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(kk.hi);
#else
ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(l_k[j ][dk_off + k]);
ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(l_k[j+1][dk_off + k]);
#endif
partial0 += dot0.s0 + dot0.s1 + dot0.s2 + dot0.s3;
partial1 += dot1.s0 + dot1.s1 + dot1.s2 + dot1.s3;
}
@ -359,7 +389,7 @@ __kernel void FA_TILE_NAME(
ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f);
FA_UNROLL
for (int k = 0; k < SPLIT_DK_VEC; k++) {
dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(l_k[j][dk_off + k]), dot_acc);
dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(FA_LK(j, dk_off + k)), dot_acc);
}
local_partial[j][tid] =
dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3;
@ -452,10 +482,21 @@ __kernel void FA_TILE_NAME(
FA_UNROLL
for (int k = 0; k < DK_VEC; k++) {
const ACC_TYPE4 qk = q_priv[k];
#if defined(FA_K_LDS_T)
// 4 KV rows adjacent in the transposed tile: two 128-bit local reads
// instead of four 64-bit ones.
const half8 kk01 = FA_LK_PAIR(k, j);
const half8 kk23 = FA_LK_PAIR(k, j + 2);
dot_acc0 = mad(qk, CONVERT_KV_ACC4(kk01.lo), dot_acc0);
dot_acc1 = mad(qk, CONVERT_KV_ACC4(kk01.hi), dot_acc1);
dot_acc2 = mad(qk, CONVERT_KV_ACC4(kk23.lo), dot_acc2);
dot_acc3 = mad(qk, CONVERT_KV_ACC4(kk23.hi), dot_acc3);
#else
dot_acc0 = mad(qk, CONVERT_KV_ACC4(l_k[j][k]), dot_acc0);
dot_acc1 = mad(qk, CONVERT_KV_ACC4(l_k[j+1][k]), dot_acc1);
dot_acc2 = mad(qk, CONVERT_KV_ACC4(l_k[j+2][k]), dot_acc2);
dot_acc3 = mad(qk, CONVERT_KV_ACC4(l_k[j+3][k]), dot_acc3);
#endif
}
ACC_TYPE s0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale;
ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale;

47
ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl
View File

@ -1631,8 +1631,25 @@ __kernel void flash_attn_f32_q4_0(
float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
#ifdef FA_HAVE_INT_DOT
// Accessors so the staging code is layout-agnostic.
#ifdef FA_K_LDS_T
#define FA_K_PACKED(ROW, IDX) l_k_packed[IDX][ROW]
#define FA_K_SCALE(ROW, BLK) l_k_scale[BLK][ROW]
#else
#define FA_K_PACKED(ROW, IDX) l_k_packed[ROW][IDX]
#define FA_K_SCALE(ROW, BLK) l_k_scale[ROW][BLK]
#endif
#ifdef FA_K_LDS_T
// K tile transposed: the 4 KV rows the QK loop walks together become adjacent, so each
// (block, group) step is ONE 128-bit local read instead of four 32-bit ones. The QK
// loop is LDS-read-issue-bound.
__local uint l_k_packed[DK_Q4_BLOCKS_PREFILL * 8][BLOCK_N];
__local float l_k_scale [DK_Q4_BLOCKS_PREFILL][BLOCK_N];
#else
__local uint l_k_packed[BLOCK_N][DK_Q4_BLOCKS_PREFILL * 8];
__local float l_k_scale [BLOCK_N][DK_Q4_BLOCKS_PREFILL];
#endif
#else
__local half4 l_k[BLOCK_N][DK_VEC];
#endif
@ -1660,17 +1677,17 @@ __kernel void flash_attn_f32_q4_0(
const global char * blk_ptr = k_base + k_row_off + blk * Q4_0_BLOCK_SIZE;
const float df = (float) vload_half(0, (const global half *) blk_ptr);
const global uchar * qs = (const global uchar *)(blk_ptr + 2);
l_k_scale[row][blk] = df;
FA_K_SCALE(row, blk) = df;
uint k_packed[8];
pack_q4_0_nibbles(qs, k_packed);
#pragma unroll
for (int j = 0; j < 8; ++j) {
l_k_packed[row][blk * 8 + j] = k_packed[j];
FA_K_PACKED(row, blk * 8 + j) = k_packed[j];
}
} else {
l_k_scale[row][blk] = 0.0f;
FA_K_SCALE(row, blk) = 0.0f;
#pragma unroll
for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u;
for (int j = 0; j < 8; ++j) FA_K_PACKED(row, blk * 8 + j) = 0u;
}
}
#else
@ -1760,6 +1777,19 @@ __kernel void flash_attn_f32_q4_0(
for (int b_local = 0; b_local < SPLIT_DK_Q4_BLOCKS; ++b_local) {
const int b = k_blk_base + b_local;
int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0;
#ifdef FA_K_LDS_T
// 4 KV rows are adjacent in the transposed tile: one 128-bit local
// read per (block, group) instead of four 32-bit ones.
#pragma unroll
for (int g = 0; g < 8; ++g) {
const uint qp = q_packed_pf[b_local * 8 + g];
const uint4 kq4 = vload4(0, &l_k_packed[b * 8 + g][j]);
sum0 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s0, sum0);
sum1 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s1, sum1);
sum2 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s2, sum2);
sum3 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s3, sum3);
}
#else
#pragma unroll
for (int g = 0; g < 8; ++g) {
const uint qp = q_packed_pf[b_local * 8 + g];
@ -1768,12 +1798,21 @@ __kernel void flash_attn_f32_q4_0(
sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2);
sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3);
}
#endif
const float qd = q_d_pf[b_local];
const int q_sum = q_sum_pf[b_local];
#ifdef FA_K_LDS_T
const float4 ks4 = vload4(0, &l_k_scale[b][j]);
s0 += (float)(sum0 - 8 * q_sum) * qd * ks4.s0;
s1 += (float)(sum1 - 8 * q_sum) * qd * ks4.s1;
s2 += (float)(sum2 - 8 * q_sum) * qd * ks4.s2;
s3 += (float)(sum3 - 8 * q_sum) * qd * ks4.s3;
#else
s0 += (float)(sum0 - 8 * q_sum) * qd * l_k_scale[j ][b];
s1 += (float)(sum1 - 8 * q_sum) * qd * l_k_scale[j+1][b];
s2 += (float)(sum2 - 8 * q_sum) * qd * l_k_scale[j+2][b];
s3 += (float)(sum3 - 8 * q_sum) * qd * l_k_scale[j+3][b];
#endif
}
#else
ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f);

53
ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl
View File

@ -1393,8 +1393,31 @@ __kernel void flash_attn_f32_q8_0(
float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
#ifdef FA_HAVE_INT_DOT
// Accessors so the staging code is layout-agnostic.
#ifdef FA_K_LDS_T
#define FA_K_PACKED(ROW, IDX) l_k_packed[IDX][ROW]
#define FA_K_SCALE(ROW, BLK) l_k_scale[BLK][ROW]
#else
#define FA_K_PACKED(ROW, IDX) l_k_packed[ROW][IDX]
#define FA_K_SCALE(ROW, BLK) l_k_scale[ROW][BLK]
#endif
#ifdef FA_K_LDS_T
// K tile transposed: [block*8 + g][kv row] instead of [kv row][block*8 + g].
//
// The QK loop walks 4 KV rows at a time against the same (b, g), so in the original
// layout those 4 values are BLOCK_N*8 uints apart and cost 4 separate 32-bit local
// reads. Transposed they are adjacent, so they are one 128-bit read -- 4x fewer LDS
// issues for the same bytes and no extra registers. That matters because the QK loop
// is LDS-read-issue-bound: a wrong-math probe that kept every dp4a but cut the LDS
// reads ran the whole kernel 41% faster (18.51 -> 10.91 ms/op), and deleting QK
// outright only reached 10.88 -- i.e. essentially ALL of QK's cost is these reads.
__local uint l_k_packed[DK_Q8_BLOCKS_PREFILL * 8][BLOCK_N];
__local float l_k_scale [DK_Q8_BLOCKS_PREFILL][BLOCK_N];
#else
__local uint l_k_packed[BLOCK_N][DK_Q8_BLOCKS_PREFILL * 8];
__local float l_k_scale [BLOCK_N][DK_Q8_BLOCKS_PREFILL];
#endif
#else
__local half4 l_k[BLOCK_N][DK_VEC];
#endif
@ -1427,7 +1450,7 @@ __kernel void flash_attn_f32_q8_0(
const global char * blk_ptr = k_base + k_row_off + blk * Q8_0_BLOCK_SIZE;
const float df = (float) vload_half(0, (const global half *) blk_ptr);
const global uchar * qs = (const global uchar *)(blk_ptr + 2);
l_k_scale[row][blk] = df;
FA_K_SCALE(row, blk) = df;
#pragma unroll
for (int j = 0; j < 8; ++j) {
uint k_packed =
@ -1435,12 +1458,12 @@ __kernel void flash_attn_f32_q8_0(
((uint) qs[j*4 + 1]) << 8 |
((uint) qs[j*4 + 2]) << 16 |
((uint) qs[j*4 + 3]) << 24;
l_k_packed[row][blk * 8 + j] = k_packed;
FA_K_PACKED(row, blk * 8 + j) = k_packed;
}
} else {
l_k_scale[row][blk] = 0.0f;
FA_K_SCALE(row, blk) = 0.0f;
#pragma unroll
for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u;
for (int j = 0; j < 8; ++j) FA_K_PACKED(row, blk * 8 + j) = 0u;
}
}
#else
@ -1556,6 +1579,19 @@ __kernel void flash_attn_f32_q8_0(
for (int b_local = 0; b_local < SPLIT_DK_Q8_BLOCKS; ++b_local) {
const int b = k_blk_base + b_local;
int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0;
#if defined(FA_K_LDS_T)
// The 4 KV rows are adjacent in the transposed tile, so each (b, g)
// step is ONE 128-bit local read instead of four 32-bit ones.
#pragma unroll
for (int g = 0; g < 8; ++g) {
const uint qp = q_packed_pf[b_local * 8 + g];
const uint4 kq4 = vload4(0, &l_k_packed[b * 8 + g][j]);
sum0 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s0, sum0);
sum1 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s1, sum1);
sum2 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s2, sum2);
sum3 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s3, sum3);
}
#else
#pragma unroll
for (int g = 0; g < 8; ++g) {
const uint qp = q_packed_pf[b_local * 8 + g];
@ -1564,11 +1600,20 @@ __kernel void flash_attn_f32_q8_0(
sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2);
sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3);
}
#endif
const float qd = q_d_pf[b_local];
#ifdef FA_K_LDS_T
const float4 ks4 = vload4(0, &l_k_scale[b][j]);
s0 += (float)sum0 * qd * ks4.s0;
s1 += (float)sum1 * qd * ks4.s1;
s2 += (float)sum2 * qd * ks4.s2;
s3 += (float)sum3 * qd * ks4.s3;
#else
s0 += (float)sum0 * qd * l_k_scale[j ][b];
s1 += (float)sum1 * qd * l_k_scale[j+1][b];
s2 += (float)sum2 * qd * l_k_scale[j+2][b];
s3 += (float)sum3 * qd * l_k_scale[j+3][b];
#endif
}
#else
ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f);