Repository: ggml-org/llama.cpp
Pull request: vulkan: fix SSM_CONV PP scaling with large ubatch sizes

Fixes #18725

The SSM_CONV shader dispatched one token per Y workgroup, each doing only `nc` (typically 4) multiply-adds. At ubatch=2048 this meant 2048 workgroups in Y with almost no work per launch — workgroup dispatch overhead dominated.

**Changes:**
- Tile 16 tokens per workgroup using a 2D local size (32×16)
- Add `vec4` dot product fast path for the common `nc=4` (d_conv) case
- Pipeline wg_denoms updated from `{32,1,1}` to `{32,16,1}`

**45/45 SSM_CONV backend-ops tests passing.**

**test-backend-ops perf (ne_a=[515,3328], nc=4):**
| | us/run | GB/s |
|--|--------|------|
| master | 526.24 | 24.29 |
| this PR | 203.56 | 62.79 |
| speedup | **2.59x** | |

**Model bench (Qwen3-Coder-Next REAM Q4_K_M, pp2048, AMD 890M):**
| ubatch | master (t/s) | this PR (t/s) | change |
|--------|-------------|--------------|--------|
| 256 | 148.43 | 150.36 | +1.3% |
| 512 | 171.38 | 175.82 | +2.6% |
| 1024 | 147.46 | 181.14 | **+22.8%** |
| 2048 | 126.34 | 161.56 | **+27.9%** |

Master shows the #18725 pattern — PP drops from 171 at ub512 to 126 at ub2048. With this fix, PP peaks at ub1024 (181) and stays strong at ub2048 (162). The degradation cliff is gone.

Tested on AMD Radeon 890M (RDNA3.5, 8 CUs, Strix Point integrated). Would appreciate testing from @lemmi on the discrete 8060S where the original issue was reported.

Changed files:
filename: ggml/src/ggml-vulkan/ggml-vulkan.cpp
patch: @@ -4576,7 +4576,7 @@ static void ggml_vk_load_shaders(vk_device& device) {
         ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d256, "ssm_scan_256_f32", ssm_scan_f32_len, ssm_scan_f32_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {256, device->subgroup_size, 16}, 1, true, true);
     }
 
-    ggml_vk_create_pipeline(device, device->pipeline_ssm_conv_f32, "ssm_conv_f32", ssm_conv_f32_len, ssm_conv_f32_data, "main", 3, sizeof(vk_op_ssm_conv_push_constants), {32, 1, 1}, {32}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_ssm_conv_f32, "ssm_conv_f32", ssm_conv_f32_len, ssm_conv_f32_data, "main", 3, sizeof(vk_op_ssm_conv_push_constants), {32, 16, 1}, {32, 16}, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_opt_step_adamw_f32, "opt_step_adamw_f32", opt_step_adamw_f32_len, opt_step_adamw_f32_data, "main", 5, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
 

filename: ggml/src/ggml-vulkan/vulkan-shaders/ssm_conv.comp
patch: @@ -5,8 +5,9 @@
 #include "types.glsl"
 
 layout(constant_id = 0) const uint BLOCK_SIZE = 32;
+layout(constant_id = 1) const uint TOKENS_PER_WG = 16;
 
-layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z = 1) in;
 
 layout(binding = 0) readonly buffer Src0 { float src0[]; };
 layout(binding = 1) readonly buffer Src1 { float src1[]; };
@@ -20,25 +21,30 @@ layout(push_constant) uniform PushConstants {
 };
 
 void main() {
-    const uint global_thread_id = gl_GlobalInvocationID.x;
-    const uint i2 = gl_WorkGroupID.y;
+    const uint i1 = gl_GlobalInvocationID.x;
+    const uint i2 = gl_WorkGroupID.y * TOKENS_PER_WG + gl_LocalInvocationID.y;
     const uint i3 = gl_WorkGroupID.z;
 
-    if (global_thread_id >= nr || i2 >= n_t || i3 >= n_s) {
+    if (i1 >= nr || i2 >= n_t || i3 >= n_s) {
         return;
     }
 
-    const uint i1 = global_thread_id;
     const uint src0_base = i3 * (nb02 / 4) + i2 + i1 * (nb01 / 4);
     const uint src1_base = i1 * (nb11 / 4);
-    const uint dst_idx = i3 * (dst_nb2 / 4) + i2 * (dst_nb1 / 4) + i1;
 
     float sum = 0.0;
-    [[unroll]] for (uint i0 = 0; i0 < nc; i0++) {
-        const uint src0_idx = src0_base + i0;
-        const uint src1_idx = src1_base + i0;
-        sum += src0[src0_idx] * src1[src1_idx];
+
+    if (nc == 4) {
+        sum = dot(
+            vec4(src0[src0_base], src0[src0_base + 1], src0[src0_base + 2], src0[src0_base + 3]),
+            vec4(src1[src1_base], src1[src1_base + 1], src1[src1_base + 2], src1[src1_base + 3])
+        );
+    } else {
+        [[unroll]] for (uint i0 = 0; i0 < nc; i0++) {
+            sum += src0[src0_base + i0] * src1[src1_base + i0];
+        }
     }
 
+    const uint dst_idx = i3 * (dst_nb2 / 4) + i2 * (dst_nb1 / 4) + i1;
     dst[dst_idx] = sum;
 }