vulkan: fix COL2IM_1D dispatch for large tensors

[AdityaAWaghmare]

Problem

The Vulkan implementation of GGML_OP_COL2IM_1D flattened the
[T_out, OC] output into a single X dispatch dimension.

For sufficiently large tensors this can exceed
maxComputeWorkGroupCount[0] on Vulkan devices with a lower X
workgroup-count limit, triggering the dispatch assertion.

This was first reproduced on Intel ADL GT2 with PocketTTS and
Qwen3-TTS for longer inputs, but the issue is not Intel-specific.
Any Vulkan device whose maxComputeWorkGroupCount[0] is smaller than
the required flattened dispatch can hit the same problem.

For example:

• T_out = 52000
• OC = 384
• total elements = 19,968,000
• local size = 256
• required X workgroups = ceil(19,968,000 / 256) = 78,000

This exceeds a device limit of 65,535.

Fix

Dispatch the two logical output dimensions separately:

• X = T_out
• Y = OC

The shader now maps gl_GlobalInvocationID.x to t_out and
gl_GlobalInvocationID.y to oc.

This changes the example dispatch from:

(78000, 1, 1)

to:

(204, 384, 1)

while preserving the existing output layout and COL2IM computation.

The existing push-constant layout is preserved.

This avoids relying on a large flattened X dispatch and makes the
operation portable across Vulkan devices with different
maxComputeWorkGroupCount limits.

Testing

Tested with Vulkan on Intel ADL GT2:

• PocketTTS, 779 characters: 11.15 s, RTF 0.261
• Qwen3-TTS 0.6B, 779 characters: 97.60 s, RTF 1.946

Both previously reproduced the maxComputeWorkGroupCount assertion.

Also tested PocketTTS on NVIDIA RTX 3050 Vulkan:

• 779 characters: 3.23 s, RTF 0.076

Qwen3-TTS on the RTX 3050 could not be tested because the 4 GB VRAM
configuration ran out of device memory.

[Copilot]

🟡 Changes recommended

Dispatch-limit handling and regression coverage remain unresolved.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Fixes Vulkan COL2IM_1D failures on large tensors by dispatching T_out and OC separately instead of flattening them into X.

Changes:

• Maps shader invocations to (t_out, oc).
• Dispatches {T_out, OC, 1} while preserving output layout and push constants.

File summaries

File	Summary
external/ggml/src/ggml-vulkan/vulkan-shaders/col2im_1d.comp	Implements 2-D indexing and bounds checks.
external/ggml/src/ggml-vulkan/ggml-vulkan.cpp	Updates dispatch dimensions. Findings: critical (1 vote)—dimensions can still exceed device limits; moderate (2 votes)—add boundary regression coverage; moderate (1 vote)—avoid excessive padding for short outputs.

Review details

Suppressed comments (1)

external/ggml/src/ggml-vulkan/ggml-vulkan.cpp:10384

• The 2-D dispatch changes the amount of work for short outputs: it launches ceil(T_out / 256) * 256 * OC invocations, versus ceil(T_out * OC / 256) * 256 before. For T_out < 256 this can be up to 256× more threads (for example, T_out = 1 launches one 256-thread workgroup per channel, with only one lane doing useful work). Consider retaining the flattened path when it is within the device X limit and selecting the 2-D path only for overflowing dispatches, or otherwise avoiding this per-row padding.

            elements = { uint32_t(dst->ne[0]), uint32_t(dst->ne[1]), 1 };

• Files reviewed: 2/2 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[AdityaAWaghmare]

A couple of points for further discussion:

The 2D dispatch approach here follows the upstream llama.cpp implementation of GGML_OP_COL2IM_1D, introduced in ggml-org/llama.cpp#24425 (commit ad39cca). The upstream implementation dispatches T_out and OC across X/Y respectively, with the same local_size_x = 256 approach used here.

Regarding the possibility of still exceeding maxComputeWorkGroupCount: that's a general property of dispatch dimensions and isn't specific to this change. The existing assertion already checks all three dispatch dimensions. The problem reproduced in this issue was specifically the flattened X dispatch: for the failing workload it required 78,000 X workgroups, exceeding the 65,535 limit on the Intel device. Splitting the logical dimensions gives (204, 384, 1) for that workload.

The concern about extra invocations for small T_out is valid and could potentially be optimized with a different/adaptive dispatch strategy. However, I kept this PR aligned with the upstream implementation and focused it on the correctness issue that caused the crash. Happy to discuss whether an adaptive approach would be preferable here.

[0xShug0]

@AdityaAWaghmare Thank you for digging deep and finding the root cause. I need more time to test it because it changes dispatch cost for every shape and have performance implications.
E.g.,

T_out=1, OC=384
old = 512 invocations
new = 98,304 invocations

It being upstream is a good sign, but we've tried bumping the GGML version before and couldn’t move forward because of performance regressions.

Could you try another apporach: keep shader untouched and cap the X dispatch size on the host side:

E.g.

case GGML_OP_COL2IM_1D:
         {
-            const uint32_t total = (uint32_t)(dst->ne[0] * dst->ne[1]);
-            elements = { total, 1, 1 };
+            const uint64_t total = uint64_t(dst->ne[0]) * uint64_t(dst->ne[1]);
+            const uint64_t max_x_elements =
+                uint64_t(ctx->device->properties.limits.maxComputeWorkGroupCount[0]) *
+                uint64_t(pipeline->wg_denoms[0]);
+
+            elements = {
+                uint32_t(std::min(total, max_x_elements)),
+                1,
+                1
+            };

The change is much smaller. The existing shader already has a grid-stride loop based on the actual dispatched workgroup count, so capping X should still cover the full output while preserving the existing dispatch for shapes that fit.

Could you test this against the failing workload and compare correctness and performance around the dispatch boundary?