| |
| /* One workgroup normalizes each row of residual = input + skip, with an |
| * optional bias. */ |
| {% if usesF16 %} |
| enable f16; |
| {% endif %} |
| {{ env.wgsl.resourceDeclarations }} |
| |
| const HIDDEN: u32 = {{ hiddenSize }}u; |
| const WG: u32 = {{ workgroupSize }}u; |
| |
| var<workgroup> partial: array<f32, WG>; |
| {% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %} |
| fn {{ name }}(value: f32, tid: u32) -> f32 { |
| {{ buffer }}[tid] = value; |
| workgroupBarrier(); |
| // Ceil-halving keeps every lane when the workgroup size is not a power of |
| // two. For even n this matches the power-of-two tree order; for odd n, lanes |
| // [0, n-half) fold the upper tail while the middle lane carries forward. |
| var n: u32 = {{ wg }}; |
| loop { |
| let half = (n + 1u) / 2u; |
| if (tid < n - half) { |
| {% if mode == "max" %} |
| {{ buffer }}[tid] = max({{ buffer }}[tid], {{ buffer }}[tid + half]); |
| {% else %} |
| {{ buffer }}[tid] = {{ buffer }}[tid] + {{ buffer }}[tid + half]; |
| {% endif %} |
| } |
| workgroupBarrier(); |
| n = half; |
| if (n == 1u) { |
| break; |
| } |
| } |
| // The default trailing barrier makes this helper safe for back-to-back calls: every lane reads |
| // slot 0 here, so the next call's first store must not run until all lanes have read it. |
| // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit. |
| let reduced = {{ buffer }}[0]; |
| {% if trailingBarrier %} |
| workgroupBarrier(); |
| {% endif %} |
| return reduced; |
| } |
| {% endmacro %} |
| |
| {{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }} |
| var<workgroup> row_inv: f32; |
| |
| fn residual_value(row: u32, d: u32) -> f32 { |
| let index = row * HIDDEN + d; |
| var value = f32(input[index]) + f32(skip[index]); |
| {% if hasBias %} |
| value = value + f32(bias[d]); |
| {% endif %} |
| return value; |
| } |
| |
| @compute @workgroup_size(WG, 1, 1) |
| fn main( |
| @builtin(workgroup_id) wg: vec3<u32>, |
| @builtin(local_invocation_id) lid: vec3<u32>) { |
| // 2D-folded row index: wg.y carries the high bits past the per-axis dispatch fold width. |
| // Reduces to wg.x when the dispatch does not fold; |
| // the row >= params.rows guard drops the over-dispatched tail. |
| let row = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u; |
| if (row >= params.rows) { |
| return; |
| } |
| let tid = lid.x; |
| |
| // RMS normalization uses one sum-of-squares sweep, without a mean or beta. |
| |
| var local_sq = 0.0; |
| for (var d: u32 = tid; d < HIDDEN; d = d + WG) { |
| let value = residual_value(row, d); |
| local_sq = local_sq + value * value; |
| } |
| let sq = reduce_sum(local_sq, tid); |
| if (tid == 0u) { |
| row_inv = inverseSqrt(sq / f32(HIDDEN) + params.epsilon); |
| } |
| workgroupBarrier(); |
| |
| for (var d: u32 = tid; d < HIDDEN; d = d + WG) { |
| let index = row * HIDDEN + d; |
| let residual = residual_value(row, d); |
| {% if writeResidualSum %} |
| input_skip_bias_sum[index] = {{ scalar }}(residual); |
| {% endif %} |
| output[index] = {{ scalar }}(residual * row_inv * f32(gamma[d])); |
| } |
| } |
| |
