metal : clamp K extent in tensor API mat-mat kernel for K not a multiple of 32 (#27450)

The Tensor API mat-mat path of kernel_mul_mm (GGML_METAL_HAS_TENSOR) fed a
static K=32 tile to the matmul2d op on every iteration. On the last, partial
K tile (ne00 % 32 != 0) the src1 slice extends past the K extent of the
tensor, and the op reads those out-of-bounds elements (undefined behavior per
the MSL specification, section 2.22.2). Depending on stale memory contents,
this corrupted the result or produced NaN.

Make the matmul2d op use dynamic_extent for K, and clamp the K extent of both
operand tensor views to the remaining valid K range (min(32, K - loop_k)) per
iteration, so the op reads exactly the valid K range on every iteration
(mirroring the tail handling of the MPP matmul2d examples). On K-aligned
inputs the clamp degenerates to the full 32-wide tile: the only difference
from the static-K op is that the dynamic-K op derives K from the operand
extents and edge-checks the tile against the tensor extents (a handful of
integer ops per iteration).

Add test-backend-ops MUL_MAT cases with K not a multiple of 32 to exercise
the unaligned K path.

Assisted-by: pi:llama.cpp/Qwen3.8-27B
This commit is contained in:
Georgi Gerganov
2026-08-20 21:31:29 +03:00
committed by GitHub
parent 6503355df0
commit a30273376e
2 changed files with 19 additions and 4 deletions
+8
View File
@@ -9298,6 +9298,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 6, 4096, 5120, {1, 1}, {1, 1}));
// K not a multiple of 32
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 64, 32, 65, {1, 1}, {1, 1}));
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 64, 32, 80, {1, 1}, {1, 1}));
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 64, 32, 80, {1, 1}, {1, 1}));
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 64, 32, 80, {1, 1}, {1, 1}));
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 64, 32, 588, {1, 1}, {1, 1})); // 14*14*3, e.g. conv_2d im2col
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 64, 32, 80, {4, 1}, {1, 1}));
#if 0
// test the mat-mat path for Metal
for (int k = 1; k < 512; ++k) {