metal : clamp K extent in tensor API mat-mat kernel for K not a multiple of 32 (llama/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
parent c16cb42f54
commit d68216aa3e
1 changed files with 11 additions and 4 deletions

View File

@ -10365,9 +10365,12 @@ kernel void kernel_mul_mm(
auto tB = tensor(ptrB, dextents<int32_t, 2>(K, N), array<int, 2>({1, strideB}));
// Configure matmul operation
// note: K is dynamic_extent (clamped to the valid range in PHASE 2), since a static
// N_MM_NK_TOTAL K tile would read src1 out of bounds when K % N_MM_NK_TOTAL != 0
// ref: https://github.com/ggml-org/llama.cpp/pull/27064
mpp::tensor_ops::matmul2d<
mpp::tensor_ops::matmul2d_descriptor(
NRB, NRA, N_MM_NK_TOTAL, false, true, true,
NRB, NRA, static_cast<int>(dynamic_extent), false, true, true,
mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
execution_simdgroups<N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y>> mm;
@ -10419,10 +10422,14 @@ kernel void kernel_mul_mm(
threadgroup_barrier(mem_flags::mem_threadgroup);
// === PHASE 2: Tensor matmul ===
auto mA = tA.slice(0, 0);
auto mB = tB.slice(loop_k, rb);
// Clamp the K extent of both operand tensors to the remaining valid K range so
// the dynamic-K op never reads past the K extent of src1 (or the staged A tile).
const int kExt = min(N_MM_NK_TOTAL, K - loop_k);
mm.run(mB, mA, cT);
auto tAv = tensor(sa, dextents<int32_t, 2>(kExt, NRA), array<int, 2>({1, N_MM_NK_TOTAL}));
auto tBv = tensor(ptrB + loop_k + rb * strideB, dextents<int32_t, 2>(kExt, N - rb), array<int, 2>({1, strideB}));
mm.run(tBv, tAv, cT);
threadgroup_barrier(mem_flags::mem_threadgroup);
}