vulkan: fuse SSM_CONV + BIAS + SILU (#22653)
This commit is contained in:
@@ -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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user