mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-11 07:19:03 +02:00
revert load reordering and scale pre-loading
This commit is contained in:
@@ -350,61 +350,33 @@ void main() {
|
||||
USE_MAGIC_BIAS ? ACC_BIAS_BITS : 0);
|
||||
}
|
||||
|
||||
const uint K_SUB = BK / TK;
|
||||
coopmat<int8_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> all_a[cms_per_row * K_SUB];
|
||||
coopmat<int8_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> all_b[cms_per_col * K_SUB];
|
||||
coopmat<int8_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> cache_a;
|
||||
coopmat<int8_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> cache_b;
|
||||
|
||||
[[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
|
||||
[[unroll]] for (uint h = 0; h < K_SUB; h++) {
|
||||
coopMatLoad(all_a[cm_row * K_SUB + h], buf_a_qs, (warp_r * WM + cm_row * TM) * QPITCH + ks * (BK / 4) + h * (TK / 4), QPITCH, gl_CooperativeMatrixLayoutRowMajor);
|
||||
}
|
||||
}
|
||||
[[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
|
||||
[[unroll]] for (uint h = 0; h < K_SUB; h++) {
|
||||
coopMatLoad(all_b[cm_col * K_SUB + h], buf_b_qs, (warp_c * WN + cm_col * TN) * QPITCH + ks * (BK / 4) + h * (TK / 4), QPITCH, gl_CooperativeMatrixLayoutColumnMajor);
|
||||
}
|
||||
}
|
||||
[[unroll]] for (uint i = 0; i < BK; i += TK) {
|
||||
[[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
|
||||
coopMatLoad(cache_a, buf_a_qs, (warp_r * WM + cm_row * TM) * QPITCH + ks * (BK / 4) + i / 4, QPITCH, gl_CooperativeMatrixLayoutRowMajor);
|
||||
|
||||
[[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
|
||||
[[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
|
||||
[[unroll]] for (uint h = 0; h < K_SUB; h++) {
|
||||
cm_result[cm_col * cms_per_row + cm_row] = coopMatMulAdd(all_a[cm_row * K_SUB + h], all_b[cm_col * K_SUB + h], cm_result[cm_col * cms_per_row + cm_row]);
|
||||
[[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
|
||||
coopMatLoad(cache_b, buf_b_qs, (warp_c * WN + cm_col * TN) * QPITCH + ks * (BK / 4) + i / 4, QPITCH, gl_CooperativeMatrixLayoutColumnMajor);
|
||||
|
||||
cm_result[cm_col * cms_per_row + cm_row] = coopMatMulAdd(cache_a, cache_b, cm_result[cm_col * cms_per_row + cm_row]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Pre-load scales into registers
|
||||
float scale_a[cms_per_row * CM_ELEMS];
|
||||
float nbias_a[cms_per_row * CM_ELEMS];
|
||||
float scale_b[cms_per_col * CM_ELEMS];
|
||||
[[unroll]] for (uint r = 0; r < cms_per_row; r++) {
|
||||
[[unroll]] for (uint e = 0; e < CM_ELEMS; e++) {
|
||||
scale_a[r * CM_ELEMS + e] = buf_a_d[ks * BM + warp_r * WM + r * TM + elem_row[e]];
|
||||
if (USE_MAGIC_BIAS) {
|
||||
nbias_a[r * CM_ELEMS + e] = -ACC_BIAS_F * scale_a[r * CM_ELEMS + e];
|
||||
}
|
||||
}
|
||||
}
|
||||
[[unroll]] for (uint c = 0; c < cms_per_col; c++) {
|
||||
[[unroll]] for (uint e = 0; e < CM_ELEMS; e++) {
|
||||
scale_b[c * CM_ELEMS + e] = buf_b_d[ks * BN + warp_c * WN + c * TN + elem_col[e]];
|
||||
}
|
||||
}
|
||||
|
||||
// Apply scales from registers
|
||||
// Apply scales directly from coopmat elements
|
||||
[[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
|
||||
[[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
|
||||
const uint tile_idx = cm_col * cms_per_row + cm_row;
|
||||
[[unroll]] for (uint e = 0; e < CM_ELEMS; e++) {
|
||||
const float da = buf_a_d[ks * BM + warp_r * WM + cm_row * TM + elem_row[e]];
|
||||
const float db = buf_b_d[ks * BN + warp_c * WN + cm_col * TN + elem_col[e]];
|
||||
if (USE_MAGIC_BIAS) {
|
||||
const float t = fma(intBitsToFloat(int(cm_result[tile_idx][e])),
|
||||
scale_a[cm_row * CM_ELEMS + e],
|
||||
nbias_a[cm_row * CM_ELEMS + e]);
|
||||
sums[tile_idx * CM_ELEMS + e] = ACC_TYPE(fma(t, scale_b[cm_col * CM_ELEMS + e],
|
||||
float(sums[tile_idx * CM_ELEMS + e])));
|
||||
const float t = fma(intBitsToFloat(int(cm_result[tile_idx][e])), da, -ACC_BIAS_F * da);
|
||||
sums[tile_idx * CM_ELEMS + e] = ACC_TYPE(fma(t, db, float(sums[tile_idx * CM_ELEMS + e])));
|
||||
} else {
|
||||
sums[tile_idx * CM_ELEMS + e] += ACC_TYPE(float(cm_result[tile_idx][e])
|
||||
* scale_a[cm_row * CM_ELEMS + e] * scale_b[cm_col * CM_ELEMS + e]);
|
||||
sums[tile_idx * CM_ELEMS + e] += ACC_TYPE(float(cm_result[tile_idx][e]) * da * db);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user