vulkan: Optimize SSM_SCAN

[jeffbolznv]

I'll leave some inline comments motivating some of the changes.

before:

Z:\github\jeffbolznv\llama.cpp\build\bin\RelWithDebInfo>llama-bench.exe -fa 1 -n 128 -p 512 -r 30 --prio 1 -m c:\models\granite-4.0-h-tiny-Q8_0.gguf
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = NVIDIA GeForce RTX 5090 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 1 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: NV_coopmat2
| model                          |       size |     params | backend    | ngl | fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | -: | --------------: | -------------------: |
| granitehybrid ?B Q8_0          |   6.88 GiB |     6.94 B | Vulkan     |  99 |  1 |           pp512 |     7110.83 ± 689.95 |
| granitehybrid ?B Q8_0          |   6.88 GiB |     6.94 B | Vulkan     |  99 |  1 |           tg128 |        193.79 ± 1.00 |

Z:\github\jeffbolznv\llama.cpp\build\bin\RelWithDebInfo>llama-bench.exe -fa 1 -n 128 -p 512 -r 30 --prio 1 -m c:\models\grante-4.0-h-tiny-Q8_0.gguf
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = NVIDIA GeForce RTX 4070 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 1 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: NV_coopmat2
| model                          |       size |     params | backend    | ngl | fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | -: | --------------: | -------------------: |
| granitehybrid ?B Q8_0          |   6.88 GiB |     6.94 B | Vulkan     |  99 |  1 |           pp512 |      3862.00 ± 19.50 |
| granitehybrid ?B Q8_0          |   6.88 GiB |     6.94 B | Vulkan     |  99 |  1 |           tg128 |        135.66 ± 0.80 |

after:

Z:\github\jeffbolznv\llama.cpp\build\bin\RelWithDebInfo>llama-bench.exe -fa 1 -n 128 -p 512 -r 30 --prio 1 -m c:\models\granite-4.0-h-tiny-Q8_0.gguf
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = NVIDIA GeForce RTX 5090 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 1 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: NV_coopmat2
| model                          |       size |     params | backend    | ngl | fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | -: | --------------: | -------------------: |
| granitehybrid ?B Q8_0          |   6.88 GiB |     6.94 B | Vulkan     |  99 |  1 |           pp512 |     9851.65 ± 131.68 |
| granitehybrid ?B Q8_0          |   6.88 GiB |     6.94 B | Vulkan     |  99 |  1 |           tg128 |        193.69 ± 1.00 |

Z:\github\jeffbolznv\llama.cpp\build\bin\RelWithDebInfo>llama-bench.exe -fa 1 -n 128 -p 512 -r 30 --prio 1 -m c:\models\granite-4.0-h-tiny-Q8_0.gguf
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = NVIDIA GeForce RTX 4070 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 1 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: NV_coopmat2
| model                          |       size |     params | backend    | ngl | fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | -: | --------------: | -------------------: |
| granitehybrid ?B Q8_0          |   6.88 GiB |     6.94 B | Vulkan     |  99 |  1 |           pp512 |      4488.07 ± 63.72 |
| granitehybrid ?B Q8_0          |   6.88 GiB |     6.94 B | Vulkan     |  99 |  1 |           tg128 |        136.61 ± 0.88 |

CC @giuseppe

[jeffbolznv]

Shifted the values of w down by a factor of 2, rather than using w>>1 everywhere. This seemed to be causing some weird code to be generated.

[jeffbolznv]

This was one too many iterations most of the time, leading to extra subgroup ops or barriers. But when D_STATE / SUBGROUP_SIZE is greater than SPLIT_H, we need at least one iteration.

[giuseppe]

thanks for fixing this one!

[jeffbolznv]

This max_idx comparison should fold away and avoid the need for the branch most of the time.

[giuseppe]

just curious, is this needed to get the subgroup to run the same code?

[giuseppe]

LGTM

I had started working on the subgroupAdd optimization, but I got stuck trying to get it to work on the Intel Arc GPU (I’m not yet sure why it was behaving differently on a NVIDIA GPU). I even suspected there might be an issue with the driver. Your version works well there too, so the issue was definitely in my version. So thanks for beating me at it :-) I will take the chance to compare the two versions and understand what I was doing wrong.

[0cc4m]

@giuseppe The usual subgroup issue with Intel would be about subgroup size (which can be between 8 and 32 for Intel) and forcing full subgroups. You probably need to do the latter to get the functions to work, and if you expect a specific size in the shader, force it. For performance on Intel 16 has been best in my experience.

[0cc4m]

LGTM