Skip to content

b10321

Choose a tag to compare

@github-actions github-actions released this 07 Aug 19:07
a194a75

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

Website:

macOS/iOS:

Linux:

Android:

Windows:

openEuler:

  • DISABLED
  • openEuler x86 (310p)
  • openEuler x86 (910b, ACL Graph)
  • openEuler aarch64 (310p)
  • openEuler aarch64 (910b, ACL Graph)

UI: