[Web] WebGPU SkipLayerNormalization / SkipSimplifiedLayerNormalization return zeros past maxStorageBufferBindingSize

<!doctype html>
<meta charset="utf-8">
<title>Norm ops past maxStorageBufferBindingSize</title>
<pre id="log">running...</pre>
<script type="module">
// Runs four norm ops on a [13200, 2560] fp32 tensor (135,168,000 bytes, just over the 128 MiB default
// maxStorageBufferBindingSize) and checks every row against a JS reference.
//   ?ort=<npm version>       onnxruntime-web version (default 1.31.1)
//   ?limits=default|adapter  device with the WebGPU default limits (default) or the adapter's maximums
const q = new URLSearchParams(location.search);
const version = q.get("ort") || "1.31.1";
const useAdapterLimits = q.get("limits") === "adapter";
const dist = `https://cdn.jsdelivr.net/npm/onnxruntime-web@${version}/dist/`;
const ort = await import(dist + "ort.webgpu.min.mjs");
ort.env.wasm.wasmPaths = dist;

const MODELS = {
 "SkipSimplifiedLayerNormalization": "CAo60wEKVAoBWAoBUwoFZ2FtbWESAVkiIFNraXBTaW1wbGlmaWVkTGF5ZXJOb3JtYWxpemF0aW9uKhEKB2Vwc2lsb24VvTeGNaABAToNY29tLm1pY3Jvc29mdBIgU2tpcFNpbXBsaWZpZWRMYXllck5vcm1hbGl6YXRpb25aFQoBWBIQCg4IARIKCgMSAU0KAwiAFFoVCgFTEhAKDggBEgoKAxIBTQoDCIAUWhQKBWdhbW1hEgsKCQgBEgUKAwiAFGIVCgFZEhAKDggBEgoKAxIBTQoDCIAUQgQKABAVQhEKDWNvbS5taWNyb3NvZnQQAQ==",
 "SkipLayerNormalization": "CAo62gEKUAoBWAoBUwoFZ2FtbWEKBGJldGESAVkiFlNraXBMYXllck5vcm1hbGl6YXRpb24qEQoHZXBzaWxvbhW9N4Y1oAEBOg1jb20ubWljcm9zb2Z0EhZTa2lwTGF5ZXJOb3JtYWxpemF0aW9uWhUKAVgSEAoOCAESCgoDEgFNCgMIgBRaFQoBUxIQCg4IARIKCgMSAU0KAwiAFFoUCgVnYW1tYRILCgkIARIFCgMIgBRaEwoEYmV0YRILCgkIARIFCgMIgBRiFQoBWRIQCg4IARIKCgMSAU0KAwiAFEIECgAQFUIRCg1jb20ubWljcm9zb2Z0EAE=",
 "SimplifiedLayerNormalization": "CAo6ugEKVgoBWAoFZ2FtbWESAVkiHFNpbXBsaWZpZWRMYXllck5vcm1hbGl6YXRpb24qFAoEYXhpcxj///////////8BoAECKhEKB2Vwc2lsb24VvTeGNaABAToAEhxTaW1wbGlmaWVkTGF5ZXJOb3JtYWxpemF0aW9uWhUKAVgSEAoOCAESCgoDEgFNCgMIgBRaFAoFZ2FtbWESCwoJCAESBQoDCIAUYhUKAVkSEAoOCAESCgoDEgFNCgMIgBRCBAoAEBVCEQoNY29tLm1pY3Jvc29mdBAB",
 "LayerNormalization": "CAo6wQEKUgoBWAoFZ2FtbWEKBGJldGESAVkiEkxheWVyTm9ybWFsaXphdGlvbioUCgRheGlzGP///////////wGgAQIqEQoHZXBzaWxvbhW9N4Y1oAEBOgASEkxheWVyTm9ybWFsaXphdGlvbloVCgFYEhAKDggBEgoKAxIBTQoDCIAUWhQKBWdhbW1hEgsKCQgBEgUKAwiAFFoTCgRiZXRhEgsKCQgBEgUKAwiAFGIVCgFZEhAKDggBEgoKAxIBTQoDCIAUQgQKABAVQhEKDWNvbS5taWNyb3NvZnQQAQ=="
};
const H = 2560, M = Number(q.get("rows")) || 13200, EPS = 1e-6;

const adapter = await navigator.gpu.requestAdapter();
const requiredLimits = useAdapterLimits
  ? { maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize, maxBufferSize: adapter.limits.maxBufferSize }
  : {};
const device = await adapter.requestDevice({
  requiredFeatures: adapter.features.has("shader-f16") ? ["shader-f16"] : [],
  requiredLimits,
});
const gpuErrors = [];
device.addEventListener("uncapturederror", (e) => gpuErrors.push(e.error.message));

const fill = (n, f) => { const a = new Float32Array(n); for (let i = 0; i < n; i++) a[i] = f(i); return a; };
const x = fill(M * H, (i) => ((i * 7919) % 1000) / 500 - 1);
const s = fill(M * H, (i) => ((i * 104729) % 1000) / 1000 - 0.5);
const gamma = fill(H, (j) => 1 + (j % 5) * 0.1);
const beta = fill(H, (j) => (j % 3) * 0.1);

function reference(op, row) {
  const skip = op.startsWith("Skip"), simplified = op.includes("Simplified");
  const v = new Float64Array(H);
  for (let j = 0; j < H; j++) v[j] = x[row * H + j] + (skip ? s[row * H + j] : 0);
  let mean = 0;
  if (!simplified) { for (const t of v) mean += t; mean /= H; }
  let sq = 0;
  for (const t of v) sq += (t - mean) ** 2;
  const inv = 1 / Math.sqrt(sq / H + EPS);
  return Array.from(v, (t, j) => (t - mean) * inv * gamma[j] + (simplified ? 0 : beta[j]));
}

const rows = [];
for (const [op, b64] of Object.entries(MODELS)) {
  const bytes = Uint8Array.from(atob(b64), (c) => c.charCodeAt(0));
  const session = await ort.InferenceSession.create(bytes, { executionProviders: [{ name: "webgpu", device }] });
  const feeds = {
    X: new ort.Tensor("float32", x, [M, H]), S: new ort.Tensor("float32", s, [M, H]),
    gamma: new ort.Tensor("float32", gamma, [H]), beta: new ort.Tensor("float32", beta, [H]),
  };
  for (const k of Object.keys(feeds)) if (!session.inputNames.includes(k)) delete feeds[k];
  const y = (await session.run(feeds)).Y.data;
  let wrongRows = 0, zeroRows = 0, firstWrongRow = -1;
  for (let r = 0; r < M; r++) {
    const want = reference(op, r);
    let wrong = false, zero = true;
    for (let j = 0; j < H; j++) {
      const got = y[r * H + j];
      if (got !== 0) zero = false;
      if (Math.abs(got - want[j]) > 1e-3 * Math.max(1, Math.abs(want[j]))) wrong = true;
    }
    if (wrong) { wrongRows++; if (firstWrongRow < 0) firstWrongRow = r; }
    if (zero) zeroRows++;
  }
  rows.push({ op, wrongRows, allZeroRows: zeroRows, firstWrongRow,
    firstWrongByte: firstWrongRow < 0 ? null : firstWrongRow * H * 4 });
  await session.release();
}

const result = {
  ort: ort.env.versions.web, userAgent: navigator.userAgent,
  adapter: adapter.info ? `${adapter.info.vendor} ${adapter.info.architecture}` : "",
  maxStorageBufferBindingSize: device.limits.maxStorageBufferBindingSize, maxBufferSize: device.limits.maxBufferSize,
  tensor: { shape: [M, H], bytes: M * H * 4 }, rows, gpuErrors,
};
document.getElementById("log").textContent = JSON.stringify(result, null, 2);
window.result = result;
</script>

# Builds the four one-node models used by repro.html and prints them as base64 (all tensors are graph inputs).
import base64, json, onnx
from onnx import helper, TensorProto as T

H = 2560
def model(op, inputs, domain="", **attrs):
    shapes = {"X": ["M", H], "S": ["M", H], "gamma": [H], "beta": [H]}
    node = helper.make_node(op, inputs, ["Y"], domain=domain, **attrs)
    g = helper.make_graph([node], op, [helper.make_tensor_value_info(n, T.FLOAT, shapes[n]) for n in inputs],
                          [helper.make_tensor_value_info("Y", T.FLOAT, ["M", H])])
    m = helper.make_model(g, opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)])
    m.ir_version = 10
    if op != "SimplifiedLayerNormalization":  # an onnxruntime op in the default domain, unknown to the checker
        onnx.checker.check_model(m)
    return base64.b64encode(m.SerializeToString()).decode()

print(json.dumps({
    "SkipSimplifiedLayerNormalization": model("SkipSimplifiedLayerNormalization", ["X", "S", "gamma"], "com.microsoft", epsilon=1e-6),
    "SkipLayerNormalization": model("SkipLayerNormalization", ["X", "S", "gamma", "beta"], "com.microsoft", epsilon=1e-6),
    "SimplifiedLayerNormalization": model("SimplifiedLayerNormalization", ["X", "gamma"], axis=-1, epsilon=1e-6),
    "LayerNormalization": model("LayerNormalization", ["X", "gamma", "beta"], axis=-1, epsilon=1e-6),
}, indent=1))