Skip to content

opencl: transpose the FA prefill K tile in local memory for perf optimization #26428

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lhez merged 1 commit into from
Aug 10, 2026
Merged

opencl: transpose the FA prefill K tile in local memory for perf optimization #26428

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions ggml/src/ggml-opencl/ggml-opencl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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;
}

Expand Down
51 changes: 46 additions & 5 deletions ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
47 changes: 43 additions & 4 deletions ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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];
Expand All @@ -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);
Expand Down
53 changes: 49 additions & 4 deletions ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -1427,20 +1450,20 @@ __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 =
(uint) qs[j*4 + 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
Expand Down Expand Up @@ -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];
Expand All @@ -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);
Expand Down
Loading