Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Hongqiang Wang
2026-08-10 11:09:19 -07:00
committed by GitHub
parent 0666ad2b2b
commit 689e227db4
4 changed files with 156 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl
+49 -4
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);