vulkan: fuse SSM_CONV + BIAS + SILU (#22653)

This commit is contained in:
Jeff Bolz
2026-05-17 10:25:50 +02:00
committed by GitHub
parent 1a68ec9378
commit 3fbadb06dc
2 changed files with 129 additions and 9 deletions
@@ -6,12 +6,15 @@
layout(constant_id = 0) const uint BLOCK_SIZE = 32;
layout(constant_id = 1) const uint TOKENS_PER_WG = 16;
layout(constant_id = 2) const bool APPLY_BIAS = false;
layout(constant_id = 3) const bool APPLY_SILU = false;
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[]; };
layout(binding = 2) buffer Dst { float dst[]; };
layout(binding = 2) readonly buffer Bias { float bias[]; };
layout(binding = 3) buffer Dst { float dst[]; };
layout(push_constant) uniform PushConstants {
uint nb01; uint nb02;
@@ -45,6 +48,13 @@ void main() {
}
}
if (APPLY_BIAS) {
sum += bias[i1];
}
if (APPLY_SILU) {
sum = sum / (1.0f + exp(-sum));
}
const uint dst_idx = i3 * (dst_nb2 / 4) + i2 * (dst_nb1 / 4) + i1;
dst[dst_idx] = sum;
}