//------------------------------------------------------------------------------
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);
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;
}
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)
#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) {
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;
}
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;
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;
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
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
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];
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);
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
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 + 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
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];
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);
overview / pkg / ggml / sources / llama.cpp / commitdiff