vulkan : fuse UNARY(GELU|SIGMOID|SILU|SOFTPLUS) + MUL - #27220
Conversation
github-actions
Bot
added
Vulkan
|
Hi @Ankk98, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
|
I'd like this to be done more consistently with existing fusions - add the code to the first shader in the sequence (unary), specialize the shaders (spec constant or ifdef) rather than runtime branching on push constants, and only fuse adjacent nodes (sort in graph_optimize only, don't try to find non-consecutive nodes at runtime, it won't work right). |
- implement fusion in unary.comp behind UNARY_MUL_FUSION ifdef, specialized pipelines per op instead of runtime branching - fuse adjacent nodes only, ordering handled by graph_optimize - drop runtime consumer scan and pending_unary_mul deferral
| }); | ||
| } | ||
|
|
||
| static bool ggml_vk_should_fuse_unary_mul(const ggml_tensor * unary, const ggml_tensor * mul) { |
There was a problem hiding this comment.
This should be grouped with the other ggml_vk_can_fuse functions (and should be called ggml_vk_can_fuse...)
| #undef CREATE_UNARY | ||
|
|
||
| #define CREATE_UNARY_MUL(name) \ | ||
| ggml_vk_create_pipeline(device, device->pipeline_ ## name ## _mul[0], #name "_mul_f32", name ## _mul_f32_len, name ## _mul_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {1}, 1); \ |
There was a problem hiding this comment.
These have norepeat=true, but seems like mul operations often use repeat. So I think we should support both.
| if (used[k]) { | ||
| continue; | ||
| } | ||
| if (std::find(current_set.begin(), current_set.end(), k) != current_set.end()) { |
There was a problem hiding this comment.
I didn't follow all of this logic, and it's concerning that it's different from some existing cases. I wonder if this could just be a small new case alongside the group of reorderings that start with "// When we've found RMS_NORM + MUL, try to find a ROPE that uses it"
1. GELU: gelu_mul_f32/f16 pipelines registered, CREATE_UNARY_MUL(gelu), GELU in dispatch + fuse gate + perf fusion name
2. Renamed/moved: gate is now ggml_vk_can_fuse_unary_mul(cgraph, unary_idx, mul_idx), placed with the other can-fuse helpers
3. norepeat both variants: each op gets plain (spec {0}) + _norepeat (spec {1}) pipelines from the same SPIR-V, selected via ggml_are_same_shape(src0, src1); the shape gate now allows broadcast (other dims equal-or-1)
4. graph_optimize: lambda deleted; standard "// UNARY + MUL: pull the consuming MUL forward" block added alongside the SSM_CONV/ROPE/MUL_MAT reorderings, with the same "other src must be weights or already processed" readiness check
… test tolerance - schedule the fused kernel like mul.comp (256 threads x 2 unrolled iterations), recovering a 10-18% prompt-processing regression - allow 5e-7 f32 error for gelu_mul: the shader evaluates gelu with an exp-based tanh identity while the CPU reference uses tanhf (~1 ulp)
| if (!ggml_is_contiguous_1(other) || !ggml_is_contiguous_1(unary->src[0])) { | ||
| return false; | ||
| } | ||
| for (int i = 0; i < GGML_MAX_DIMS; ++i) { |
There was a problem hiding this comment.
I think this could just use ggml_can_repeat?
Ankk98
changed the title
The fused kernel indexes src1 via per-dim fastmod (generic_binary_head.glsl), which is exact whenever the other operand tiles into the unary result -- not just when its dims are equal or 1. Replace the hand-rolled loop with ggml_can_repeat(other, unary) so the check matches the kernel's actual capability and reuses the standard helper. Argument order matters: reversed, it would wrongly admit graphs where the unary result is mul->src[1] and the other operand is larger, producing truncated output. Also add a rep_ne0 layout to the fused unary+mul backend tests covering a non-1 repeat factor along dim 0.
gemma4's per-layer embedding gating builds gelu -> view_2d_slice -> mul, where the intervening view is a zero-compute node aliasing an input that was computed much earlier. Strict adjacency requirements meant neither CUDA nor the vulkan unary+mul fusion handled this pattern. Extend ggml_vk_graph_optimize to detect a UNARY whose consuming MUL is separated only by unscheduled zero-compute nodes (GGML_OP_NONE, VIEW, RESHAPE, TRANSPOSE, PERMUTE) and schedule those nodes ahead of the pair, making it adjacent so the existing fusion applies. The reorder is guarded by ggml_vk_can_fuse_unary_mul, a source-availability check for every interleaved node, and the protected fusion patterns (topk_moe*, snake); if fusion is later rejected the reordered graph still executes correctly, just unfused. Add a view_mid layout to the fused unary+mul backend tests replicating the gemma4 pattern.
| continue; | ||
| } | ||
|
|
||
| // UNARY + EMPTY* + MUL: when only zero-compute nodes (GGML_OP_NONE, |
There was a problem hiding this comment.
This block of code seems overcomplicated and unnecessary. The code at line 17836 should be able to pull the MUL ahead of a view.
| get_indices(idx, i00, i01, i02, i03); | ||
|
|
||
| data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = | ||
| D_TYPE(FLOAT_TYPE(OP(float(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]))) * FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)])); |
There was a problem hiding this comment.
Can we also support running the OP on the B operand? I think this happens in qwen models:
// Apply sigmoid to the gate
shared_gate = ggml_sigmoid(ctx0, shared_gate);
cb(shared_gate, "shared_expert_gate_sigmoid", il);
// Apply the gate to the shared expert output
ffn_shexp = ggml_mul(ctx0, ffn_shexp, shared_gate);
2 participants
Overview
Issue #27194
Additional information
Design:
op/dtype, repeat + norepeat variants selected by shape equality.
the standard ggml_can_repeat(other, unary) broadcast check.
(NONE/VIEW/RESHAPE/TRANSPOSE/PERMUTE) out of UNARY→MUL gaps so view-separated sites fuse
(per-layer embedding gating in gemma4/gemma3n builds this shape).
Tests:
Muse-Glimmer-30B, gemma4-12B.
Requirements