Spaces:

---

webml-community

/

gemma-4-webgpu-kernels

like 294

Running

Thanks for the report! Can you please provide more info about your OS/browser/hardware? That will be helpful. Thanks!

Version 149.0.7827.116 (Official Build) (64-bit)
dell g16 i9 13th gen 32GB RAM rtx 4070 - 8GB vRAM
all latest cuda and vulkan cores installed
windows latest updated version 11

thanks for that info! If you're able to, maybe you can run this diagnostic script too: https://huggingface.co/spaces/webml-community/gemma-4-webgpu-kernels/discussions/1#6a340073bac7944cc4c62476

I'm trying to fix it on another nvidia laptop of mine; hope to get it done soon!

I've just submitted a PR to fix this, hosting the fixed version here: https://huggingface.co/spaces/igorls/gemma-4-webgpu-kernels

same issue, different text, every time ive tried. win11, 4090, chrome

same issue, different text, every time ive tried. win11, 4090, chrome

try this version: https://huggingface.co/spaces/igorls/gemma-4-webgpu-kernels

@igorls
fixed properly well done
"
Gemma-4 WebGPU Subgroup & Softmax Loop Investigation Report

This report provides a deep-dive analysis of how the WebGPU hardware incompatibility and attention softmax loop issues were resolved in the gemma-4-webgpu-kernels environment.

1. Executive Summary & Root Cause Analysis

On-device LLMs utilizing WebGPU (like Gemma-4 E2B) rely heavily on subgroup operations (also known as SIMD/wavefront operations) inside the attention softmax and normalization (RMSNorm) kernels. Subgroups allow threads in the same execution warp to share data directly via high-speed register shuffles (subgroupShuffleXor) without writing to slower shared workgroup memory.

The original implementation compiled by Fable 5 / Google was optimized exclusively for Apple M-Series hardware, introducing two critical bugs when running on Windows/Linux with NVIDIA, AMD, or Intel hardware:

Issue A: The AMD Softmax Collapse (Repetitive Loops)

The Math Bug: The original WGSL shaders hardcoded a 5-step butterfly reduction loop (1u, 2u, 4u, 8u, 16u), which assumes a physical hardware subgroup size of exactly 32.

AMD Behavior: AMD GPUs run with a physical wavefront/subgroup size of 64 (Wave64). Because the reduction stopped shuffling at 16u, the upper 32 lanes of the subgroup never communicated with the lower 32 lanes.

The Failure: The attention matrix was calculated with unreduced, incomplete query-key maximums. This caused the subsequent softmax layer to experience numerical overflow/collapse, pinning the model's activations to high-frequency tokens (resulting in endless repetitions of "the the the..." or "Aula Aula...").

Issue B: The Intel RMSNorm Validation Crash

The Math Bug: Intel integrated and discrete GPUs often execute WebGPU with physical subgroup sizes of 8 or 16.

Intel Behavior: The original shader compiled the instruction subgroupShuffleXor(x, 16u).

The Failure: Under the WebGPU shading specification (WGSL), a shuffle offset must be strictly less than the physical subgroup size. Shuffling at offset 16 on an 8-lane or 16-lane subgroup is an out-of-bounds compilation error. This caused the browser's shader compiler to instantly invalidate the RMSNorm module, crashing the initialization.

2. The Mechanics of the "Igorls" Solution

The solved repository resolves these issues by combining two powerful strategies: synchronous compilation interception and safe fallback detection.

Strategy 1: Compiler Interception (API Hooking)

WebGPU GPUShaderModule objects are opaque in the browser; their compiled WGSL source code cannot be read or edited after creation. To solve this, the patch overrides GPUDevice.prototype.createShaderModule to inspect and intercept the raw text strings of the shader descriptor before they are sent to the GPU driver.

// Intercepting the shader creation API
const originalCreateShaderModule = GPUDevice.prototype.createShaderModule;

GPUDevice.prototype.createShaderModule = function(descriptor) {
if (descriptor && descriptor.code && typeof descriptor.code === 'string') {
let code = descriptor.code;

    // Match subgroup shuffles that hardcode a 32-lane limit
    if (code.includes('subgroupShuffleXor') && code.includes('16u') && !code.includes('subgroup_size')) {
        // Apply regular expression to modify WGSL source in memory
        descriptor.code = patchWGSLReduction(code);
    }
}
return originalCreateShaderModule.call(this, descriptor);

};

Strategy 2: Dynamic, Hardware-Aware Butterfly Reductions

Instead of assuming a static subgroup size of 32, the compiler patch uses a type-agnostic regular expression to capture the minified variable name of the reduction parameters and replaces the hardcoded shuffles with a hardware-aware loop using WebGPU's @builtin(subgroup_size):

The Patched WGSL Code

// Before (Hardcoded to assume subgroup_size == 32)
fn sg_sum(v: f32) -> f32 {
var x = v;
x = x + subgroupShuffleXor(x, 1u);
x = x + subgroupShuffleXor(x, 2u);
x = x + subgroupShuffleXor(x, 4u);
x = x + subgroupShuffleXor(x, 8u);
x = x + subgroupShuffleXor(x, 16u); // Out of bounds on Intel!
return x;
}

// After (Type-Agnostic & Bound-Safe)
fn sg_sum(v: f32) -> f32 {
var val_patched = v;
let ssz = subgroup_size; // Query actual physical subgroup size at runtime

val_patched = val_patched + subgroupShuffleXor(val_patched, 1u);
val_patched = val_patched + subgroupShuffleXor(val_patched, 2u);
val_patched = val_patched + subgroupShuffleXor(val_patched, 4u);
val_patched = val_patched + subgroupShuffleXor(val_patched, 8u);

// Dynamically scale shuffle steps based on actual hardware bounds
if (ssz >= 32u) { 
    val_patched = val_patched + subgroupShuffleXor(val_patched, 16u); 
}
if (ssz >= 64u) { 
    val_patched = val_patched + subgroupShuffleXor(val_patched, 32u); 
}
if (ssz >= 128u) { 
    val_patched = val_patched + subgroupShuffleXor(val_patched, 64u); 
}
return val_patched;

}

This prevents Intel compilation crashes because offsets > ssz are hidden behind hardware branches.

This fixes the AMD Wave64 softmax loop because shuffles scale up to 32u (covering all 64 lanes), compiling perfect reduction trees.

Strategy 3: Subgroups Feature Deselection (Fallback Force)

On certain GPU architectures where subgroup control flow can behave unpredictably depending on driver versions, the most robust way to ensure stability is to trigger the model's Workgroup Shared Memory fallback.

The original gemma-4-e2b.js file has two built-in shader paths:

Subgroups Path: Fast, used if 'subgroups' is present in device.features.

Workgroup Memory Path: Highly stable fallback, used if 'subgroups' is absent.

The patch hooks requestDevice and intercepts GPUAdapter.prototype.features.has to hide the 'subgroups' capability string on non-Apple hardware (like Windows AMD/Intel):

// Hook features checking
if (window.GPUSupportedFeatures && GPUSupportedFeatures.prototype.has) {
const originalHas = GPUSupportedFeatures.prototype.has;
GPUSupportedFeatures.prototype.has = function(feature) {
if (feature === 'subgroups' && !isAppleGPU) {
return false; // Force hide subgroups
}
return originalHas.call(this, feature);
};
}

By stripping this capability at initialization, the engine drops subgroup shuffles entirely and falls back to workgroup shared arrays, eliminating hardware-specific math bugs while maintaining maximum execution compatibility.

3. Summary of Impact

Metric / Fix

Original Behavior

Patched Behavior (igorls Solution)

NVIDIA Compatibility

Softmax Loop (Gibberish Output)

Fully coherent attention states, flawless generation.

AMD Compatibility

Softmax Loop (Wave64 Divergence)

Coherent outputs via full 64-lane unrolling or workgroup fallback.

Intel Compatibility

Instant RMSNorm compiler crash

Flawless compilation and loading.

Precision Support

Mixed (f32 hardcoded signatures)

Fully type-agnostic (supports f16, QAT, and f32 models).
"