llama.cpp
a194a75b - metal : fix NORM/RMS_NORM for row lengths that leave a partial simdgroup (#26708)

Commit
2 days ago
metal : fix NORM/RMS_NORM for row lengths that leave a partial simdgroup (#26708) ggml_metal_op_norm sized the threadgroup with `nth = std::min(nth, args.ne00_t)`, which can leave nth not a multiple of the simdgroup size. The kernels finish their row reduction with a cross-simdgroup step where each lane of the last simdgroup reads one per-simdgroup partial sum out of shmem_f32: if (tiisg == 0) { shmem_f32[sgitg] = sumf; } threadgroup_barrier(mem_flags::mem_threadgroup); sumf = shmem_f32[tiisg]; sumf = simd_sum(sumf); When the last simdgroup is partial it has fewer lanes than the threadgroup has simdgroups, so the tail of the partial sums is never read and the row sum is too small. For ne00_t = 33 nth becomes 33: two simdgroups, but only one lane in the second, so one of the two partial sums is dropped. The mean and variance are then wrong for the whole row. Round ne00_t up to a whole number of simdgroups instead. Rounding up rather than dropping the clamp keeps the threadgroup as small as possible: deleting the line would raise nth to the next power of two (ne00_t = 544 -> 1024 instead of 544), which costs idle lanes on 26 row lengths below 8192 that were already correct, including 1536 and 3584. GGML_OP_NORM is affected as well as GGML_OP_RMS_NORM - both dispatch through ggml_metal_op_norm. No mainstream LLM hidden size hits this: ne00_t is ne00/4 on the vectorized path, so 4096, 8192, 2048 and friends all give a multiple of 32. It is reachable from other norm shapes, e.g. 320-channel norms. Add NORM and RMS_NORM cases for ne0 = 33, 132 and 260 across the existing eps values. 33 exercises the scalar path and 132/260 the vectorized one, since only those divide by 4. Before, on M3 Pro: test-backend-ops test -b MTL0 -o NORM 25/50 test-backend-ops test -b MTL0 -o RMS_NORM 26/51 After: test-backend-ops test -b MTL0 -o NORM 50/50 test-backend-ops test -b MTL0 -o RMS_NORM 51/51 test-backend-ops test -b MTL0 13943/13943
Author
Parents
Loading