Skip to content

Commit 587f602

Browse files
shalinib-ibmtaronaeo
authored andcommitted
ggml-cpu: support K tails in power10 Q8/Q4 MMA matmul (llama/24753)
* ggml-cpu: support K tails in Power10 MMA Q8/Q4 matmul This patch removes the requirement that K be divisible by kc in the tinyBlas_Q0_PPC tiled matmul path. Process the final K panel using its actual depth and pass the reduced panel size through packing and kernel execution. This allows more workloads to use the MMA kernel and reduces fallback to mnpack. * Apply suggestion from @taronaeo Co-authored-by: Aaron Teo <taronaeo@gmail.com> --------- Co-authored-by: Aaron Teo <taronaeo@gmail.com>
1 parent 50fd261 commit 587f602

1 file changed

Lines changed: 6 additions & 5 deletions

File tree

ggml/src/ggml-cpu/llamafile/sgemm.cpp

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2345,7 +2345,7 @@ class tinyBLAS_Q0_PPC {
23452345
else if (n_aligned % 16 == 0) nc = 16;
23462346
else nc = 8;
23472347
}
2348-
bool can_use_tiled = n_aligned > 0 && (m % mc == 0) && (k % kc == 0);
2348+
bool can_use_tiled = n_aligned > 0 && (m % mc == 0);
23492349
if (can_use_tiled) {
23502350
matmul_tiled(m, n_aligned, mc, nc, kc);
23512351
if (n > n_aligned) {
@@ -3063,13 +3063,14 @@ class tinyBLAS_Q0_PPC {
30633063
int64_t ii = (job / xtiles) * mc;
30643064
int64_t jj = (job % xtiles) * nc;
30653065
for (int64_t kk = 0; kk < k; kk += kc) {
3066+
int64_t k_cur = MIN(kc, k - kk);
30663067
if constexpr(is_Ablock_q4) {
3067-
packNormal_q4_fp16(A + ii * lda + kk, lda, mc, kc, (uint8_t *)A_pack);
3068+
packNormal_q4_fp16(A + ii * lda + kk, lda, mc, k_cur, (uint8_t *)A_pack);
30683069
} else {
3069-
packNormal_q8_fp16(A + ii * lda + kk, lda, mc, kc, (uint8_t *)A_pack);
3070+
packNormal_q8_fp16(A + ii * lda + kk, lda, mc, k_cur, (uint8_t *)A_pack);
30703071
}
3071-
packNormal_q8_fp16(B + jj * ldb + kk, ldb, nc, kc, (uint8_t *)B_pack);
3072-
KERNEL_Q0(ii, jj, mc, nc, kc, kk, A_pack, B_pack);
3072+
packNormal_q8_fp16(B + jj * ldb + kk, ldb, nc, k_cur, (uint8_t *)B_pack);
3073+
KERNEL_Q0(ii, jj, mc, nc, k_cur, kk, A_pack, B_pack);
30733074
}
30743075
}
30753076
}

0 commit comments

Comments
 (0)