CANN: fix acl_rstd allocation size in ggml_cann_rms_norm

[noemotiovon]

Fixes #15330

Adjust the allocation size of acl_rstd. The parameter dims is set to 3 according to the CANN documentation.

Co-authored-by: Yuchuan [email protected]

[noemotiovon]

@yuchuan-cao, could you review my code?

[noemotiovon]

Opt Test:

ggml_backend_cann_context: device 0 async operator submission is OFF
ggml_backend_cann_context: device 0 execution mode is GRAPH (acl graph enabled)
Backend 1/2: CANN0
  Device description: Ascend910B4
  Device memory: 30196 MB (29848 MB free)

new_pool_for_device: device 0 use vmm pool
  RMS_NORM(type=f32,ne=[64,5,4,3],v=0,eps=0.000000): OK
  RMS_NORM(type=f32,ne=[64,5,4,3],v=1,eps=0.000000): OK
  RMS_NORM(type=f32,ne=[64,5,4,3],v=0,eps=0.000001): OK
  RMS_NORM(type=f32,ne=[64,5,4,3],v=1,eps=0.000001): OK
  RMS_NORM(type=f32,ne=[64,5,4,3],v=0,eps=0.000100): OK
  RMS_NORM(type=f32,ne=[64,5,4,3],v=1,eps=0.000100): OK
  RMS_NORM(type=f32,ne=[64,5,4,3],v=0,eps=0.100000): OK
  RMS_NORM(type=f32,ne=[64,5,4,3],v=1,eps=0.100000): OK
  11837/11837 tests passed
  Backend CANN0: OK
Backend 2/2: CPU
  Skipping
2/2 backends passed
OK

[hipudding]

You can jsut use src as rstd_ne, acl_rstd_ne array is not necessary.

[noemotiovon]

That approach won’t work. If we simply reuse src->ne, then for a 3D tensor the resulting dimensions would be {ne0, ne1, ne2}, whereas what we actually need are {ne1, ne2, ne3}.

[yuchuan-cao]

I suppose preparing src->ne sized buffer would work fine, and preparing acl_rstd_ne sized buffer would be a more memory-efficient solution.

I suggest you double-check the changes in llama-bench. The confusing thing is that the memory problem in #15330 doesn't occur in test-backend-ops and llama-cli. It only emerges when running llama-bench. I failed to find out why.

[noemotiovon]

@yuchuan-cao, Thanks for the suggestion. I ran the llama-bench test, and it works fine.

# script
./bin/llama-bench -fa 0 -n 0 -p 512 -r 50 -m /home/lichenguang25/.ollama/models/blobs/sha256-6f96e01a3f550ca08aea1e5725bb8d5a7eccc6f281c30417e9d380b8c46467bd -ngl 99

# log
| model                          |       size |     params | backend    | ngl |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --------------: | -------------------: |
| qwen2 1B F16                   | 942.43 MiB |   494.03 M | CANN       |  99 |           pp512 |      3261.30 ± 18.48 |