luboslenco
46 lines
1.4 KiB
Plaintext
46 lines
1.4 KiB
Plaintext
#version 450
|
|
/*
|
|
* Qwen3 RMSNorm: per-row normalization over `hidden`, scaled by a bf16 weight.
|
|
* out[s,i] = x[s,i] / sqrt(mean(x[s,:]^2) + eps) * weight[i]
|
|
* One workgroup per row; the row's elements are reduced in shared memory.
|
|
* Activations (in/out) are f32; the weight is bf16 (packed two-per-uint).
|
|
*/
|
|
layout(local_size_x = 256) in;
|
|
|
|
layout(std430, binding = 0) readonly buffer In { float x[]; };
|
|
layout(std430, binding = 1) readonly buffer W { uint wbf[]; };
|
|
layout(std430, binding = 2) writeonly buffer Out { float outp[]; };
|
|
|
|
layout(push_constant) uniform PC { uint rows, hidden; float eps; } pc;
|
|
|
|
float load_bf16(uint e) {
|
|
uint w = wbf[e >> 1u];
|
|
uint h = ((e & 1u) == 0u) ? (w & 0xffffu) : (w >> 16u);
|
|
return uintBitsToFloat(h << 16u);
|
|
}
|
|
|
|
shared float s_sq[256];
|
|
|
|
void main() {
|
|
uint row = gl_WorkGroupID.x;
|
|
uint tid = gl_LocalInvocationID.x;
|
|
uint base = row * pc.hidden;
|
|
|
|
float sq = 0.0;
|
|
for (uint i = tid; i < pc.hidden; i += 256u) {
|
|
float v = x[base + i];
|
|
sq += v * v;
|
|
}
|
|
s_sq[tid] = sq;
|
|
barrier();
|
|
for (uint s = 128u; s > 0u; s >>= 1u) {
|
|
if (tid < s) s_sq[tid] += s_sq[tid + s];
|
|
barrier();
|
|
}
|
|
|
|
float inv = inversesqrt(s_sq[0] / float(pc.hidden) + pc.eps);
|
|
for (uint i = tid; i < pc.hidden; i += 256u) {
|
|
outp[base + i] = x[base + i] * inv * load_bf16(i);
|
|
}
|
|
}
|