vulkan: Add State Space Model (SSM) Operations Support #16463

giuseppe · 2025-10-07T13:28:48Z

implement SSM scan and SSM conv for Vulkan.

Tested on an NVIDIA L4:
master (4e0388a)

# bin/llama-bench -r 5 -m ~/models/granite-4.0-h-tiny-UD-Q8_K_XL.gguf 

| model                          |       size |     params | backend    | ngl |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --------------: | -------------------: |
| granitehybrid ?B Q8_0          |   7.73 GiB |     6.94 B | Vulkan     | 999 |           pp512 |       1419.59 ± 6.62 |
| granitehybrid ?B Q8_0          |   7.73 GiB |     6.94 B | Vulkan     | 999 |           tg128 |         33.57 ± 0.19 |

and db8b8bc:

# bin/llama-bench -r 5 -m ~/models/granite-4.0-h-tiny-UD-Q8_K_XL.gguf

| model                          |       size |     params | backend    | ngl |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --------------: | -------------------: |
| granitehybrid ?B Q8_0          |   7.73 GiB |     6.94 B | Vulkan     |  99 |       4 |     512 |     4096 |    0 |           pp512 |      2569.27 ± 19.37 |
| granitehybrid ?B Q8_0          |   7.73 GiB |     6.94 B | Vulkan     |  99 |       4 |     512 |     4096 |    0 |           tg128 |         80.77 ± 0.12 |

jeffbolznv

Thanks for this contribution!

ggml/src/ggml-vulkan/ggml-vulkan.cpp

jeffbolznv · 2025-10-07T14:35:10Z

ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp

+    warp_sdata[warp_offset + lane] = val;
+    barrier();
+
+    if (lane < 16) warp_sdata[warp_offset + lane] += warp_sdata[warp_offset + lane + 16];


This seems like it's assuming a subgroup size of 32 (also at line 37).

Do I understand correctly that this doesn't actually rely on a subgroup size of 32, but it's splitting the workgroup into groups of 32 and just reducing those (and it looks like some reduction across groups of 32 has already happened?).

sorry I've missed this one. Yeah I don't think it would work with a size != 32. I need to think more through this one.

Do you've any suggestions on what I could do here?

I think this may work because you're not relying on SubgroupInvocationId or SubgroupID, you've just split the workgroup into groups of 32. Maybe we can just test it on AMD (with wave64) and Intel and verify that it works.

it work on Intel, but I am worried about all these settings that we made configurable. I've not really tried how it behaves with different values of the constants we defined. Or is the assumption that these values should not be tweaked from vulkan-shaders-gen.cpp without also changing the implementation in the shader?

You could change this to a while loop that will handle any power-of-two value of WARP_SIZE. We do want to allow the spec constants to be changeable but it's fine to have limitations like "must be a power of two".

Hmm wave64 AMD and wave8 llvmpipe are failing one test here, possibly due to this. All other tests are passing.

[SSM_SCAN] NMSE = 31335529439335960.000000000 > 0.000000100 SSM_SCAN(type=f32,d_state=16,head_dim=1,n_head=1024,n_group=1,n_seq_tokens=32,n_seqs=4): FAIL

I think Intel also has a subgroup size of 32 so it wouldn't be a good test for this.

ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp

ggml/src/ggml-vulkan/ggml-vulkan.cpp

ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp

jeffbolznv · 2025-10-07T14:50:00Z

ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp

+    return warp_sdata[warp_offset];
+}
+
+void main() {


Do all threads always load/store in bounds? In the host code there was some rounding up going on, which suggests maybe some threads don't correspond to in-bounds locations.

Ping on this one. I don't really understand what this shader does and which locations it should be accessing.

I've tried to follow what the CUDA shader does. I'll spend more time on it and see if there is anything I can improve about memory access and make sure all the assumptions in the code are checked.

ggml/src/ggml-vulkan/vulkan-shaders/ssm_conv.comp

giuseppe · 2025-10-08T10:23:47Z

I've addressed the comments and pushed a new version. The results are even better now:

ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = NVIDIA L4 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 49152 | int dot: 0 | matrix cores: KHR_coopmat
| model                          |       size |     params | backend    | ngl | threads | n_batch | n_ubatch | mmap |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | ------: | ------: | -------: | ---: | --------------: | -------------------: |
| granitehybrid ?B Q8_0          |   7.73 GiB |     6.94 B | Vulkan     |  99 |       4 |     512 |     4096 |    0 |           pp512 |      2569.27 ± 19.37 |
| granitehybrid ?B Q8_0          |   7.73 GiB |     6.94 B | Vulkan     |  99 |       4 |     512 |     4096 |    0 |           tg128 |         80.77 ± 0.12 |

ggml/src/ggml-vulkan/ggml-vulkan.cpp

jeffbolznv · 2025-10-08T15:11:10Z

ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp


+    string_to_spv("ssm_scan_f32_d16", "ssm_scan.comp", {{"A_TYPE", "float"}});
+    string_to_spv("ssm_scan_f32_d128", "ssm_scan.comp", {{"A_TYPE", "float"}});
+    string_to_spv("ssm_scan_f32_d256", "ssm_scan.comp", {{"A_TYPE", "float"}});


These three are all identical now, you only need one.

jeffbolznv · 2025-10-08T15:12:37Z

ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp

+    warp_sdata[warp_offset + lane] = val;
+    barrier();
+
+    if (lane < 16) warp_sdata[warp_offset + lane] += warp_sdata[warp_offset + lane + 16];


You could change this to a while loop that will handle any power-of-two value of WARP_SIZE. We do want to allow the spec constants to be changeable but it's fine to have limitations like "must be a power of two".

giuseppe · 2025-10-09T09:57:35Z

I've completely replaced the code to reduce sum with subgroupAdd() as sum.comp does. It didn't work earlier because I had to use gl_SubgroupInvocationID == 0 instead of % WARP_SIZE == 0.

In the last version I've also renamed WARP_SIZE to SUBGROUP_SIZE.

0cc4m · 2025-10-09T10:00:30Z

Be aware that not all devices support subgroup commands. If there's a performance advantage to using them, you can do that, but it would still need a fallback to using a shared memory reduction. If there isn't a performance advantage, just use the shared memory reduction for compatibility.

giuseppe · 2025-10-09T10:18:04Z

Be aware that not all devices support subgroup commands. If there's a performance advantage to using them, you can do that, but it would still need a fallback to using a shared memory reduction. If there isn't a performance advantage, just use the shared memory reduction for compatibility.

with the subgroup code I get:

ggml_vulkan: 0 = NVIDIA L4 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 49152 | int dot: 0 | matrix cores: KHR_coopmat
| model                          |       size |     params | backend    | ngl | threads | n_batch | n_ubatch | mmap |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | ------: | ------: | -------: | ---: | --------------: | -------------------: |
| granitehybrid ?B Q8_0          |   7.73 GiB |     6.94 B | Vulkan     |  99 |       4 |     512 |     4096 |    0 |           pp512 |      2692.95 ± 49.85 |
| granitehybrid ?B Q8_0          |   7.73 GiB |     6.94 B | Vulkan     |  99 |       4 |     512 |     4096 |    0 |           tg128 |         81.13 ± 0.08 |

if I revert to the reduction loop, I've:

ggml_vulkan: 0 = NVIDIA L4 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 49152 | int dot: 0 | matrix cores: KHR_coopmat
| model                          |       size |     params | backend    | ngl | threads | n_batch | n_ubatch | mmap |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | ------: | ------: | -------: | ---: | --------------: | -------------------: |
| granitehybrid ?B Q8_0          |   7.73 GiB |     6.94 B | Vulkan     |  99 |       4 |     512 |     4096 |    0 |           pp512 |      2562.04 ± 31.72 |
| granitehybrid ?B Q8_0          |   7.73 GiB |     6.94 B | Vulkan     |  99 |       4 |     512 |     4096 |    0 |           tg128 |         81.14 ± 0.09 |

giuseppe · 2025-10-09T10:23:55Z

I've reverted to the version with a for loop. We can look at the subgroup optimization later

Add State Space Model scan operation to the Vulkan backend. Signed-off-by: Giuseppe Scrivano <[email protected]>

Add State Space Model conv operation to the Vulkan backend. Signed-off-by: Giuseppe Scrivano <[email protected]>

github-actions bot added Vulkan Issues specific to the Vulkan backend ggml changes relating to the ggml tensor library for machine learning labels Oct 7, 2025

giuseppe mentioned this pull request Oct 7, 2025

Misc. bug: Slower performance on newer models Apriel-1.5-15b, granite-4.0-h-tiny #16454

Open

giuseppe marked this pull request as ready for review October 7, 2025 14:05

giuseppe requested a review from 0cc4m as a code owner October 7, 2025 14:05

jeffbolznv requested changes Oct 7, 2025

View reviewed changes

giuseppe force-pushed the ssm-vulkan branch from db8b8bc to 3a58b42 Compare October 8, 2025 10:23

0cc4m reviewed Oct 8, 2025

View reviewed changes

ggml/src/ggml-vulkan/ggml-vulkan.cpp Show resolved Hide resolved

giuseppe force-pushed the ssm-vulkan branch from 3a58b42 to 694f01d Compare October 8, 2025 12:54

jeffbolznv reviewed Oct 8, 2025

View reviewed changes

giuseppe force-pushed the ssm-vulkan branch from 694f01d to c53c551 Compare October 9, 2025 09:44

giuseppe force-pushed the ssm-vulkan branch from c53c551 to a03a716 Compare October 9, 2025 10:23

giuseppe force-pushed the ssm-vulkan branch 3 times, most recently from 742ebd7 to c467631 Compare October 9, 2025 12:04

giuseppe added 2 commits October 9, 2025 14:26

vulkan: implement SSM scan operation

f72338e

Add State Space Model scan operation to the Vulkan backend. Signed-off-by: Giuseppe Scrivano <[email protected]>

vulkan: implement SSM conv operation

6e70718

Add State Space Model conv operation to the Vulkan backend. Signed-off-by: Giuseppe Scrivano <[email protected]>

giuseppe force-pushed the ssm-vulkan branch from c467631 to 6e70718 Compare October 9, 2025 12:28

vulkan: Add State Space Model (SSM) Operations Support #16463

Are you sure you want to change the base?

vulkan: Add State Space Model (SSM) Operations Support #16463

Conversation

giuseppe commented Oct 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jeffbolznv left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

netrunnereve Oct 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

giuseppe commented Oct 8, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

giuseppe commented Oct 9, 2025

Uh oh!

0cc4m commented Oct 9, 2025

Uh oh!

giuseppe commented Oct 9, 2025

Uh oh!

giuseppe commented Oct 9, 2025

Uh oh!

Uh oh!

giuseppe commented Oct 7, 2025 •

edited

Loading

netrunnereve Oct 8, 2025 •

edited

Loading