Upload shaders/rmsnorm.wgsl with huggingface_hub

shaders/rmsnorm.wgsl

@@ -0,0 +1,66 @@
+// rmsnorm.wgsl — Root Mean Square LayerNorm (no mean subtraction).
+// Ported from gfx1151_inference/shaders/rmsnorm.comp.
+//   y[t, d] = x[t, d] / sqrt(mean(x[t, :]^2) + eps) * weight[d]
+// One workgroup per row. Tree reduction for variance.
+
+const WG_SIZE: u32 = 64u;
+
+struct Params {
+    rows: u32,
+    D: u32,
+    eps: f32,
+}
+
+@group(0) @binding(0) var<storage, read>       x_buf: array<f32>;
+@group(0) @binding(1) var<storage, read>       w_buf: array<f32>;
+@group(0) @binding(2) var<storage, read_write> y_buf: array<f32>;
+
+@group(1) @binding(0) var<uniform> params: Params;
+
+var<workgroup> partial_sumsq: array<f32, 64>;
+
+@compute @workgroup_size(WG_SIZE)
+fn main(
+    @builtin(workgroup_id) wg_id: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>
+) {
+    let row = wg_id.x;
+    let tid = lid.x;
+    if (row >= params.rows) { return; }
+
+    let base = row * params.D;
+
+    // 1. Each thread accumulates its strided slice of x[row]^2
+    var s: f32 = 0.0;
+    var d: u32 = tid;
+    loop {
+        if (d >= params.D) { break; }
+        let v = x_buf[base + d];
+        s = s + v * v;
+        d = d + WG_SIZE;
+    }
+    partial_sumsq[tid] = s;
+    workgroupBarrier();
+
+    // 2. Tree reduction
+    var off: u32 = WG_SIZE / 2u;
+    loop {
+        if (off == 0u) { break; }
+        if (tid < off) {
+            partial_sumsq[tid] = partial_sumsq[tid] + partial_sumsq[tid + off];
+        }
+        workgroupBarrier();
+        off = off >> 1u;
+    }
+
+    let mean_sq = partial_sumsq[0u] / f32(params.D);
+    let scale = 1.0 / sqrt(mean_sq + params.eps);
+
+    // 3. Apply: y = x * scale * weight
+    d = tid;
+    loop {
+        if (d >= params.D) { break; }
+        y_buf[base + d] = x_buf[base + d] * scale * w_buf[d];
+        d = d + WG_SIZE;
+    }
+}