Spaces:

---

webml-community

/

gemma-4-webgpu-kernels

like 297

Running

This is the root-cause fix for the corrupted output on NVIDIA/Windows reported in discussions #1 and #8 (and what looks like the same symptom in #6, #10, and #11).

What is happening. Twelve reduction helpers in gemma-4-e2b.js already guard on sgExact32: hardware subgroupAdd on fixed 32-wide adapters, a 32-lane subgroupShuffleXor butterfly on ranged-width adapters. Eight sites (the QatMatMul family and its split-K siblings) emit a bare subgroupAdd unconditionally. On an RTX 5070 (D3D12, adapter range 32 to 128, Chrome 150) we captured both compiled kernel sets (subgroups on and off), ran all 59 matched pairs on identical inputs on the same device, and diffed: 53 pairs match exactly, 6 diverge, and all six are QatMatMul instantiations with N_ROWS = 2, the shapes that interleave subgroup reductions with the lane-divergent if (tid == 0u) stores. The divergence is deterministic and large (relative error up to 1.95), which is the gibberish users see.

Two details worth knowing before reviewing:

1. It is not a wave-width problem. A near-identical variant of the real kernel that exports @builtin(subgroup_size) reports 32 while still producing wrong sums. The WGSL is spec-valid; the platform miscompiles a subgroupAdd that executes right after lane-divergent control flow. A standalone 2-second reproducer (no model, no quantization, a plain 60-line float kernel) fails 20 of 20 trials on the 5070 and passes 20 of 20 with the butterfly, so the Dawn or driver teams can bisect the layer directly. On this machine Chrome's path is Tint to HLSL 2021 to DXC to the NVIDIA driver (610.62), for whoever picks that up.
2. Two independent changes each fix all six instantiations: this butterfly guard, or hoisting the reductions above the stores. The butterfly matches what the other twelve sites already do, so that is what this PR does. It is one mechanical edit repeated eight times; nothing else changes (551,802 bytes to 553,858).

Validation.

• Apple M1: sgExact32 renders the same WGSL as before, byte for byte. Loads and generates unchanged (41 tok/s in our run).
• RTX 5070: with the identical text transformation applied at shader-module creation time to the running engine, all 6 divergent pairs go exact, output is coherent, an 11-test task battery passes 11 of 11, long-context stress shows no repetition loops, and decode runs at 213 to 223 tok/s versus 193 tok/s with subgroups disabled (the #8 workaround) and 218 tok/s (corrupted) stock. Warm TTFT 0.65 s.

Full write-up, raw JSON evidence, the standalone reproducer, and a drop-in runtime guard for people running the current bundle today: https://github.com/Ar5en1c/gemma4-webgpu-nvidia-subgroup-fix

Credit to @Intellipedia (#1) for first flagging the NVIDIA adapter behavior and to @igorls (#8) for confirming the corruption on an RTX PRO 6000 and proposing the disable-subgroups workaround this supersedes.