Upload shaders/rmsnorm_noweight.wgsl with huggingface_hub

shaders/rmsnorm_noweight.wgsl

@@ -0,0 +1,53 @@
+// rmsnorm_noweight.wgsl — RMSNorm without learned weights.
+// Used for Falcon-Mamba's B, C, dt_pre normalization.
+//   y[d] = x[d] / sqrt(mean(x[:]^2) + eps)
+// Single workgroup, one row. D can be small (16 or 256).
+
+const WG_SIZE: u32 = 64u;
+
+struct Params {
+    D: u32,
+    eps: f32,
+}
+
+@group(0) @binding(0) var<storage, read_write> 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(local_invocation_id) lid: vec3<u32>) {
+    let tid = lid.x;
+
+    var s: f32 = 0.0;
+    var d: u32 = tid;
+    loop {
+        if (d >= params.D) { break; }
+        let v = buf[d];
+        s = s + v * v;
+        d = d + WG_SIZE;
+    }
+    partial_sumsq[tid] = s;
+    workgroupBarrier();
+
+    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);
+
+    d = tid;
+    loop {
+        if (d >= params.D) { break; }
+        buf[d] = buf[d] * scale;
+        d = d + WG_SIZE;
+    }
+}