{"schemaVersion":1,"id":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference","title":"WebGPU Compute Shaders for ML Inference: Client-Side Model Execution","description":"Running ML models directly in the browser with WebGPU compute shaders. WGSL implementation details, benchmarks against WebGL and server-side inference, and practical limitations of client-side model execution.","language":"en","datePublished":"2025-10-01","dateModified":null,"contentFormat":"text/markdown","content":"WebGPU\n\n# WebGPU Compute Shaders for ML Inference: Client-Side Model Execution\n\n2025-10-019 min read\n\nLucio Durán\nEngineering Manager & AI Solutions Architect\n\nAlso available in: [Español](<https://lucioduran.com/es/blog/webgpu-compute-shaders-ml-inference>)\n\n## Why Not Just Use ONNX Runtime Web?\n\n[ONNX Runtime Web](<https://onnxruntime.ai/>) is a solid tool for production use. However, it is an abstraction over WebGPU (or WebGL, or WASM), and abstractions have costs. Profiling an ONNX Runtime Web session running a DistilBERT model reveals that roughly 18% of the total inference time was spent on operator dispatch overhead, memory allocation bookkeeping, and graph optimization passes that had already been applied to the model.\n\nWriting compute shaders directly gives you control over:\n\n- Exact workgroup dimensions and dispatch geometry\n- Memory layout in storage and uniform buffers\n- Synchronization between passes\n- Quantized arithmetic without runtime dequantization\n\nThe tradeoff is that you're writing GPU code. There's no free lunch.\n\n## Setting Up the WebGPU Pipeline\n\nFirst, the boilerplate. The adapter/device negotiation code is standard boilerplate; the interesting parts follow.\n\n```javascript\nconst adapter = await navigator.gpu.requestAdapter({\n powerPreference: 'high-performance'\n});\n\nconst device = await adapter.requestDevice({\n requiredFeatures: ['timestamp-query'],\n requiredLimits: {\n maxStorageBufferBindingSize: 256 * 1024 * 1024,\n maxBufferSize: 256 * 1024 * 1024,\n maxComputeWorkgroupSizeX: 256,\n maxComputeInvocationsPerWorkgroup: 256,\n }\n});\n```\n\nThat ` maxStorageBufferBindingSize ` request is important, the default is 128MB on many implementations, and a quantized 125M parameter model can easily exceed that when you include KV cache buffers.\n\n## The Core: Matrix Multiplication in WGSL\n\nThe backbone of transformer inference is matrix multiplication. Here's a WGSL compute shader that does tiled matmul with shared memory, which is the minimum viable approach for acceptable performance:\n\n```wgsl\n@group(0) @binding(0) var<storage, read> a: array<f32>;\n@group(0) @binding(1) var<storage, read> b: array<f32>;\n@group(0) @binding(2) var<storage, read_write> result: array<f32>;\n\nstruct Dimensions {\n M: u32,\n N: u32,\n K: u32,\n}\n@group(0) @binding(3) var<uniform> dims: Dimensions;\n\nconst TILE_SIZE: u32 = 16u;\n\nvar<workgroup> tileA: array<array<f32, 16>, 16>;\nvar<workgroup> tileB: array<array<f32, 16>, 16>;\n\n@compute @workgroup_size(16, 16, 1)\nfn main(\n @builtin(global_invocation_id) gid: vec3<u32>,\n @builtin(local_invocation_id) lid: vec3<u32>,\n) {\n let row = gid.x;\n let col = gid.y;\n var sum: f32 = 0.0;\n\n let numTiles = (dims.K + TILE_SIZE - 1u) / TILE_SIZE;\n\n for (var t: u32 = 0u; t < numTiles; t = t + 1u) {\n let tiledCol = t * TILE_SIZE + lid.y;\n let tiledRow = t * TILE_SIZE + lid.x;\n\n if (row < dims.M && tiledCol < dims.K) {\n tileA[lid.x][lid.y] = a[row * dims.K + tiledCol];\n } else {\n tileA[lid.x][lid.y] = 0.0;\n }\n\n if (tiledRow < dims.K && col < dims.N) {\n tileB[lid.x][lid.y] = b[tiledRow * dims.N + col];\n } else {\n tileB[lid.x][lid.y] = 0.0;\n }\n\n workgroupBarrier();\n\n for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {\n sum = sum + tileA[lid.x][k] * tileB[k][lid.y];\n }\n\n workgroupBarrier();\n }\n\n if (row < dims.M && col < dims.N) {\n result[row * dims.N + col] = sum;\n }\n}\n```\n\nThis is a standard tiled approach. The key detail most tutorials miss: those two ` workgroupBarrier() ` calls are non-negotiable. The first ensures the tile is fully loaded before computation; the second ensures computation is complete before the tile memory is reused. Remove either one and you'll get intermittent wrong results that are absolute hell to debug.\n\n## Going Quantized: INT4 Matmul\n\nRunning FP32 matmul works, but it's wasteful for inference. Most production models use INT4 or INT8 quantization. Here's where it gets interesting, WGSL doesn't have native int4 types, so you pack two int4 values into a single u8 (well, u32, since WGSL's smallest type is 32-bit).\n\n```wgsl\nfn unpack_int4(packed: u32, index: u32) -> f32 {\n let shift = (index % 8u) * 4u;\n let nibble = (packed >> shift) & 0xFu;\n // Sign-extend from 4-bit\n let signed = i32(nibble) - select(0, 16, nibble >= 8u);\n return f32(signed);\n}\n\n@compute @workgroup_size(256, 1, 1)\nfn quantized_matmul(\n @builtin(global_invocation_id) gid: vec3<u32>,\n) {\n let row = gid.x / dims.N;\n let col = gid.x % dims.N;\n\n if (row >= dims.M) { return; }\n\n var sum: f32 = 0.0;\n let groupSize: u32 = 32u;\n\n for (var k: u32 = 0u; k < dims.K; k = k + 1u) {\n let group_idx = k / groupSize;\n let scale = scales[row * (dims.K / groupSize) + group_idx];\n let zero = zeros[row * (dims.K / groupSize) + group_idx];\n\n let packed_idx = (row * dims.K + k) / 8u;\n let w = unpack_int4(weights_packed[packed_idx], k);\n let dequant = (w - zero) * scale;\n\n sum = sum + input[col * dims.K + k] * dequant;\n }\n\n result[row * dims.N + col] = sum;\n}\n```\n\nThe ` unpack_int4 ` function is doing the heavy lifting. We store 8 int4 values per u32, unpack on the fly, and apply group-wise dequantization. The group size of 32 matches what GPTQ and AWQ use, so you can load quantized models directly.\n\n## The Attention Mechanism\n\nSelf-attention is the other performance-critical piece. A common implementation error is as follows: they implement attention as three separate dispatches (Q*K^T, softmax, attn*V). But you can fuse the softmax into the same pass as the Q\\*K^T multiply if you use the online softmax trick:\n\n```wgsl\n@compute @workgroup_size(256, 1, 1)\nfn fused_attention(\n @builtin(global_invocation_id) gid: vec3<u32>,\n @builtin(local_invocation_id) lid: vec3<u32>,\n) {\n let head = gid.x / seq_len;\n let query_pos = gid.x % seq_len;\n\n if (head >= num_heads) { return; }\n\n var max_score: f32 = -1e10;\n var sum_exp: f32 = 0.0;\n var acc: array<f32, 64>; // head_dim\n\n // Online softmax: single pass over K\n for (var k: u32 = 0u; k < seq_len; k = k + 1u) {\n var score: f32 = 0.0;\n for (var d: u32 = 0u; d < head_dim; d = d + 1u) {\n score += q[head * seq_len * head_dim + query_pos * head_dim + d]\n * keys[head * seq_len * head_dim + k * head_dim + d];\n }\n score = score / sqrt(f32(head_dim));\n\n let new_max = max(max_score, score);\n let correction = exp(max_score - new_max);\n sum_exp = sum_exp * correction + exp(score - new_max);\n\n for (var d: u32 = 0u; d < head_dim; d = d + 1u) {\n acc[d] = acc[d] * correction\n + exp(score - new_max)\n * values[head * seq_len * head_dim + k * head_dim + d];\n }\n max_score = new_max;\n }\n\n // Normalize\n for (var d: u32 = 0u; d < head_dim; d = d + 1u) {\n output[head * seq_len * head_dim + query_pos * head_dim + d]\n = acc[d] / sum_exp;\n }\n}\n```\n\nThe online softmax avoids a second pass over the attention scores to normalize. On a 512-token sequence with 12 heads, this fusion saved ~4ms per layer on an RTX 3060, which adds up fast across 6 or 12 layers.\n\n## The Dispatch Orchestration\n\nTying it all together in JavaScript. This is where you coordinate the multi-pass inference:\n\n```javascript\nclass WebGPUInferenceEngine {\n constructor(device, model) {\n this.device = device;\n this.pipelines = {};\n this.buffers = {};\n this.model = model;\n }\n\n async buildPipelines() {\n this.pipelines.embed = await this.createComputePipeline(embeddingShader);\n this.pipelines.layerNorm = await this.createComputePipeline(layerNormShader);\n this.pipelines.attention = await this.createComputePipeline(fusedAttentionShader);\n this.pipelines.ffn = await this.createComputePipeline(quantizedMatmulShader);\n this.pipelines.softmax = await this.createComputePipeline(finalSoftmaxShader);\n }\n\n runInference(inputTokens) {\n const encoder = this.device.createCommandEncoder();\n\n // Embedding lookup\n this.dispatchEmbedding(encoder, inputTokens);\n\n // Transformer layers\n for (let layer = 0; layer < this.model.numLayers; layer++) {\n this.dispatchLayerNorm(encoder, layer, 'pre_attn');\n this.dispatchAttention(encoder, layer);\n this.dispatchLayerNorm(encoder, layer, 'pre_ffn');\n this.dispatchFFN(encoder, layer);\n }\n\n // Final norm + logits\n this.dispatchLayerNorm(encoder, -1, 'final');\n this.dispatchLogits(encoder);\n\n this.device.queue.submit([encoder.finish()]);\n }\n}\n```\n\nAll the dispatches go into a single command encoder. The GPU driver handles dependency tracking between dispatches based on buffer usage, this is one area where WebGPU actually shines compared to doing this in CUDA, where you'd manually manage streams.\n\n## Benchmarks: WebGPU vs WebGL vs Server\n\nI ran these on a 2024 MacBook Pro (M3 Pro) and a desktop with an RTX 3060, using a quantized DistilBERT (66M params, INT4):\n\n| Method | M3 Pro | RTX 3060 | Notes |\n| --- | --- | --- | --- |\n| WebGPU Compute (ours) | 19ms | 23ms | Custom WGSL shaders |\n| ONNX Runtime Web (WebGPU) | 31ms | 38ms | Same model, default config |\n| ONNX Runtime Web (WebGL) | 67ms | 54ms | WebGL2 backend |\n| Server (Lambda + ONNX) | 45ms\\* | 45ms\\* | \\*Includes ~30ms network |\n| Server (Lambda cold start) | 2.8s | 2.8s | Warm-up penalty |\n\nThe M3 Pro numbers are noteworthy. Apple's WebGPU implementation is genuinely excellent, their Metal backend for WebGPU is mature and well-optimized. The RTX 3060 numbers are slightly higher because Chrome's Dawn backend adds a thin Vulkan translation layer.\n\nThe real win isn't the raw latency, it's the elimination of cold starts and network variability. A P99 latency of 25ms in the browser beats a P99 of 3 seconds from a serverless function that might need to cold-start.\n\n## Memory Layout Pitfalls\n\nA common pitfall worth highlighting: WGSL's storage buffer alignment rules are strict and poorly documented. Every struct member must be aligned to a multiple of its size, and the struct itself must be aligned to the largest member alignment. For a mixed f32/u32 struct:\n\n```wgsl\nstruct LayerParams {\n weight_offset: u32, // offset 0, align 4\n bias_offset: u32, // offset 4, align 4\n scale: f32, // offset 8, align 4\n num_elements: u32, // offset 12, align 4\n // Total: 16 bytes, align 4, this is fine\n}\n\n// But if you add a vec4:\nstruct LayerParamsV2 {\n weight_offset: u32, // offset 0\n // 12 bytes padding!\n params: vec4<f32>, // offset 16, align 16\n num_elements: u32, // offset 32\n // 12 bytes padding!\n // Total: 48 bytes\n}\n```\n\nThat padding is silent and brutal. Your JavaScript-side buffer writes will be wrong if you don't account for it. Use ` offset_of ` annotations or just avoid vec types in uniform structs.\n\n## When NOT to Use This Approach\n\nThere are important limitations to consider. This approach is not suitable when:\n\n- **Your model is larger than ~500M parameters.** GPU memory in browser contexts is limited, and you're sharing it with the renderer.\n- **You need deterministic results across devices.** Floating-point behavior varies between GPU vendors in ways that matter for some applications.\n- **Your users are on integrated graphics.** Intel UHD 630 can run these shaders, but at 10x the latency. Feature-detect and have a fallback.\n- **You need to keep model weights private.** Everything shipped to the browser is inspectable. If your model is proprietary, this isn't for you.\n\n## What's Next\n\nAn active area of exploration is WebGPU implementation of flash attention for longer context windows. The challenge is that WGSL doesn't give you the fine-grained shared memory control that CUDA does, so the tiling strategy needs to be different. This will be worth examining once meaningful benchmarks become available.\n\nThe broader point is this: the industry spent a decade assuming that ML inference belongs on servers. For a growing category of models and use cases, the client GPU is not just adequate, it's superior. Zero latency variance, zero cold starts, zero server costs. The tooling is finally catching up to make this practical.\n\n## Frequently asked questions\n\n### Can WebGPU really compete with server-side inference for production ML?\n\nFor models under ~500M parameters with proper quantization, absolutely. We measured 23ms inference on a quantized DistilBERT with WebGPU compute shaders on a mid-range laptop GPU, which is competitive with a cold-start serverless function when you factor in network latency.\n\n### What browsers support WebGPU compute shaders right now?\n\nAs of early 2026, Chrome 121+, Edge 121+, and Firefox 130+ ship with full WebGPU compute support. Safari has partial support behind a flag. For production, you should feature-detect with navigator.gpu and fall back to ONNX Runtime Web with WebGL.\n\n### How does WGSL compare to GLSL for compute workloads?\n\nWGSL is more explicit about memory layout and has better safety guarantees. You lose some of the syntactic sugar from GLSL compute shaders, but you gain predictable behavior across vendors. The type system catches buffer layout mismatches at pipeline creation time rather than at runtime.\n\n## Tips\n\n1. **Workgroup sizing matters more than you think:** On most GPUs, ` @workgroup_size(256, 1, 1) ` outperforms ` (64, 4, 1) ` for 1D reductions even though the thread count is the same. The memory coalescing pattern changes drastically with the X dimension.\n2. **Avoid dynamic indexing in WGSL tight loops:** Use ` let cached = weights[idx]; ` outside the inner loop. WGSL doesn't guarantee the compiler will hoist storage buffer reads, and you'll pay the latency on every iteration.\n3. **Double-buffer your storage bindings:** Instead of a single output buffer with a barrier, alternate between ` @group(0) @binding(0) ` and ` @group(0) @binding(1) ` across dispatch passes. This lets the GPU overlap compute and memory operations on architectures that support it.\n4. **Profile with ` GPUCommandEncoder.writeTimestamp() `:** Wrap your dispatch calls with timestamp queries on a dedicated ` GPUQuerySet ` of type ` 'timestamp' `. Resolve into a mappable buffer and read back, this gives you actual GPU-side microsecond timings, not CPU-side noise.\n\n- webgpu\n- wgsl\n- machine-learning\n- browser-inference\n- compute-shaders\n\n[Back to Blog](<https://lucioduran.com/blog>)","headings":[{"level":1,"text":"WebGPU Compute Shaders for ML Inference: Client-Side Model Execution","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#title"},{"level":2,"text":"Why Not Just Use ONNX Runtime Web?","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#why-not-just-use-onnx-runtime-web"},{"level":2,"text":"Setting Up the WebGPU Pipeline","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#setting-up-the-webgpu-pipeline"},{"level":2,"text":"The Core: Matrix Multiplication in WGSL","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#the-core-matrix-multiplication-in-wgsl"},{"level":2,"text":"Going Quantized: INT4 Matmul","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#going-quantized-int4-matmul"},{"level":2,"text":"The Attention Mechanism","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#the-attention-mechanism"},{"level":2,"text":"The Dispatch Orchestration","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#the-dispatch-orchestration"},{"level":2,"text":"Benchmarks: WebGPU vs WebGL vs Server","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#benchmarks-webgpu-vs-webgl-vs-server"},{"level":2,"text":"Memory Layout Pitfalls","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#memory-layout-pitfalls"},{"level":2,"text":"When NOT to Use This Approach","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#when-not-to-use-this-approach"},{"level":2,"text":"What's Next","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#whats-next"},{"level":2,"text":"Frequently asked questions","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#questions"},{"level":3,"text":"Can WebGPU really compete with server-side inference for production ML?","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#question-1"},{"level":3,"text":"What browsers support WebGPU compute shaders right now?","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#question-2"},{"level":3,"text":"How does WGSL compare to GLSL for compute workloads?","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#question-3"},{"level":2,"text":"Tips","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#article-tips"}],"sections":[{"heading":"WebGPU Compute Shaders for ML Inference: Client-Side Model Execution","level":1,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference","content":"WebGPU"},{"heading":"WebGPU Compute Shaders for ML Inference: Client-Side Model Execution","level":1,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#title","content":"2025-10-019 min read\n\nLucio Durán\nEngineering Manager & AI Solutions Architect\n\nAlso available in: [Español](<https://lucioduran.com/es/blog/webgpu-compute-shaders-ml-inference>)"},{"heading":"Why Not Just Use ONNX Runtime Web?","level":2,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#why-not-just-use-onnx-runtime-web","content":"[ONNX Runtime Web](<https://onnxruntime.ai/>) is a solid tool for production use. However, it is an abstraction over WebGPU (or WebGL, or WASM), and abstractions have costs. Profiling an ONNX Runtime Web session running a DistilBERT model reveals that roughly 18% of the total inference time was spent on operator dispatch overhead, memory allocation bookkeeping, and graph optimization passes that had already been applied to the model.\n\nWriting compute shaders directly gives you control over:\n\n- Exact workgroup dimensions and dispatch geometry\n- Memory layout in storage and uniform buffers\n- Synchronization between passes\n- Quantized arithmetic without runtime dequantization\n\nThe tradeoff is that you're writing GPU code. There's no free lunch."},{"heading":"Setting Up the WebGPU Pipeline","level":2,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#setting-up-the-webgpu-pipeline","content":"First, the boilerplate. The adapter/device negotiation code is standard boilerplate; the interesting parts follow.\n\n```javascript\nconst adapter = await navigator.gpu.requestAdapter({\n powerPreference: 'high-performance'\n});\n\nconst device = await adapter.requestDevice({\n requiredFeatures: ['timestamp-query'],\n requiredLimits: {\n maxStorageBufferBindingSize: 256 * 1024 * 1024,\n maxBufferSize: 256 * 1024 * 1024,\n maxComputeWorkgroupSizeX: 256,\n maxComputeInvocationsPerWorkgroup: 256,\n }\n});\n```\n\nThat ` maxStorageBufferBindingSize ` request is important, the default is 128MB on many implementations, and a quantized 125M parameter model can easily exceed that when you include KV cache buffers."},{"heading":"The Core: Matrix Multiplication in WGSL","level":2,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#the-core-matrix-multiplication-in-wgsl","content":"The backbone of transformer inference is matrix multiplication. Here's a WGSL compute shader that does tiled matmul with shared memory, which is the minimum viable approach for acceptable performance:\n\n```wgsl\n@group(0) @binding(0) var<storage, read> a: array<f32>;\n@group(0) @binding(1) var<storage, read> b: array<f32>;\n@group(0) @binding(2) var<storage, read_write> result: array<f32>;\n\nstruct Dimensions {\n M: u32,\n N: u32,\n K: u32,\n}\n@group(0) @binding(3) var<uniform> dims: Dimensions;\n\nconst TILE_SIZE: u32 = 16u;\n\nvar<workgroup> tileA: array<array<f32, 16>, 16>;\nvar<workgroup> tileB: array<array<f32, 16>, 16>;\n\n@compute @workgroup_size(16, 16, 1)\nfn main(\n @builtin(global_invocation_id) gid: vec3<u32>,\n @builtin(local_invocation_id) lid: vec3<u32>,\n) {\n let row = gid.x;\n let col = gid.y;\n var sum: f32 = 0.0;\n\n let numTiles = (dims.K + TILE_SIZE - 1u) / TILE_SIZE;\n\n for (var t: u32 = 0u; t < numTiles; t = t + 1u) {\n let tiledCol = t * TILE_SIZE + lid.y;\n let tiledRow = t * TILE_SIZE + lid.x;\n\n if (row < dims.M && tiledCol < dims.K) {\n tileA[lid.x][lid.y] = a[row * dims.K + tiledCol];\n } else {\n tileA[lid.x][lid.y] = 0.0;\n }\n\n if (tiledRow < dims.K && col < dims.N) {\n tileB[lid.x][lid.y] = b[tiledRow * dims.N + col];\n } else {\n tileB[lid.x][lid.y] = 0.0;\n }\n\n workgroupBarrier();\n\n for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {\n sum = sum + tileA[lid.x][k] * tileB[k][lid.y];\n }\n\n workgroupBarrier();\n }\n\n if (row < dims.M && col < dims.N) {\n result[row * dims.N + col] = sum;\n }\n}\n```\n\nThis is a standard tiled approach. The key detail most tutorials miss: those two ` workgroupBarrier() ` calls are non-negotiable. The first ensures the tile is fully loaded before computation; the second ensures computation is complete before the tile memory is reused. Remove either one and you'll get intermittent wrong results that are absolute hell to debug."},{"heading":"Going Quantized: INT4 Matmul","level":2,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#going-quantized-int4-matmul","content":"Running FP32 matmul works, but it's wasteful for inference. Most production models use INT4 or INT8 quantization. Here's where it gets interesting, WGSL doesn't have native int4 types, so you pack two int4 values into a single u8 (well, u32, since WGSL's smallest type is 32-bit).\n\n```wgsl\nfn unpack_int4(packed: u32, index: u32) -> f32 {\n let shift = (index % 8u) * 4u;\n let nibble = (packed >> shift) & 0xFu;\n // Sign-extend from 4-bit\n let signed = i32(nibble) - select(0, 16, nibble >= 8u);\n return f32(signed);\n}\n\n@compute @workgroup_size(256, 1, 1)\nfn quantized_matmul(\n @builtin(global_invocation_id) gid: vec3<u32>,\n) {\n let row = gid.x / dims.N;\n let col = gid.x % dims.N;\n\n if (row >= dims.M) { return; }\n\n var sum: f32 = 0.0;\n let groupSize: u32 = 32u;\n\n for (var k: u32 = 0u; k < dims.K; k = k + 1u) {\n let group_idx = k / groupSize;\n let scale = scales[row * (dims.K / groupSize) + group_idx];\n let zero = zeros[row * (dims.K / groupSize) + group_idx];\n\n let packed_idx = (row * dims.K + k) / 8u;\n let w = unpack_int4(weights_packed[packed_idx], k);\n let dequant = (w - zero) * scale;\n\n sum = sum + input[col * dims.K + k] * dequant;\n }\n\n result[row * dims.N + col] = sum;\n}\n```\n\nThe ` unpack_int4 ` function is doing the heavy lifting. We store 8 int4 values per u32, unpack on the fly, and apply group-wise dequantization. The group size of 32 matches what GPTQ and AWQ use, so you can load quantized models directly."},{"heading":"The Attention Mechanism","level":2,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#the-attention-mechanism","content":"Self-attention is the other performance-critical piece. A common implementation error is as follows: they implement attention as three separate dispatches (Q*K^T, softmax, attn*V). But you can fuse the softmax into the same pass as the Q\\*K^T multiply if you use the online softmax trick:\n\n```wgsl\n@compute @workgroup_size(256, 1, 1)\nfn fused_attention(\n @builtin(global_invocation_id) gid: vec3<u32>,\n @builtin(local_invocation_id) lid: vec3<u32>,\n) {\n let head = gid.x / seq_len;\n let query_pos = gid.x % seq_len;\n\n if (head >= num_heads) { return; }\n\n var max_score: f32 = -1e10;\n var sum_exp: f32 = 0.0;\n var acc: array<f32, 64>; // head_dim\n\n // Online softmax: single pass over K\n for (var k: u32 = 0u; k < seq_len; k = k + 1u) {\n var score: f32 = 0.0;\n for (var d: u32 = 0u; d < head_dim; d = d + 1u) {\n score += q[head * seq_len * head_dim + query_pos * head_dim + d]\n * keys[head * seq_len * head_dim + k * head_dim + d];\n }\n score = score / sqrt(f32(head_dim));\n\n let new_max = max(max_score, score);\n let correction = exp(max_score - new_max);\n sum_exp = sum_exp * correction + exp(score - new_max);\n\n for (var d: u32 = 0u; d < head_dim; d = d + 1u) {\n acc[d] = acc[d] * correction\n + exp(score - new_max)\n * values[head * seq_len * head_dim + k * head_dim + d];\n }\n max_score = new_max;\n }\n\n // Normalize\n for (var d: u32 = 0u; d < head_dim; d = d + 1u) {\n output[head * seq_len * head_dim + query_pos * head_dim + d]\n = acc[d] / sum_exp;\n }\n}\n```\n\nThe online softmax avoids a second pass over the attention scores to normalize. On a 512-token sequence with 12 heads, this fusion saved ~4ms per layer on an RTX 3060, which adds up fast across 6 or 12 layers."},{"heading":"The Dispatch Orchestration","level":2,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#the-dispatch-orchestration","content":"Tying it all together in JavaScript. This is where you coordinate the multi-pass inference:\n\n```javascript\nclass WebGPUInferenceEngine {\n constructor(device, model) {\n this.device = device;\n this.pipelines = {};\n this.buffers = {};\n this.model = model;\n }\n\n async buildPipelines() {\n this.pipelines.embed = await this.createComputePipeline(embeddingShader);\n this.pipelines.layerNorm = await this.createComputePipeline(layerNormShader);\n this.pipelines.attention = await this.createComputePipeline(fusedAttentionShader);\n this.pipelines.ffn = await this.createComputePipeline(quantizedMatmulShader);\n this.pipelines.softmax = await this.createComputePipeline(finalSoftmaxShader);\n }\n\n runInference(inputTokens) {\n const encoder = this.device.createCommandEncoder();\n\n // Embedding lookup\n this.dispatchEmbedding(encoder, inputTokens);\n\n // Transformer layers\n for (let layer = 0; layer < this.model.numLayers; layer++) {\n this.dispatchLayerNorm(encoder, layer, 'pre_attn');\n this.dispatchAttention(encoder, layer);\n this.dispatchLayerNorm(encoder, layer, 'pre_ffn');\n this.dispatchFFN(encoder, layer);\n }\n\n // Final norm + logits\n this.dispatchLayerNorm(encoder, -1, 'final');\n this.dispatchLogits(encoder);\n\n this.device.queue.submit([encoder.finish()]);\n }\n}\n```\n\nAll the dispatches go into a single command encoder. The GPU driver handles dependency tracking between dispatches based on buffer usage, this is one area where WebGPU actually shines compared to doing this in CUDA, where you'd manually manage streams."},{"heading":"Benchmarks: WebGPU vs WebGL vs Server","level":2,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#benchmarks-webgpu-vs-webgl-vs-server","content":"I ran these on a 2024 MacBook Pro (M3 Pro) and a desktop with an RTX 3060, using a quantized DistilBERT (66M params, INT4):\n\n| Method | M3 Pro | RTX 3060 | Notes |\n| --- | --- | --- | --- |\n| WebGPU Compute (ours) | 19ms | 23ms | Custom WGSL shaders |\n| ONNX Runtime Web (WebGPU) | 31ms | 38ms | Same model, default config |\n| ONNX Runtime Web (WebGL) | 67ms | 54ms | WebGL2 backend |\n| Server (Lambda + ONNX) | 45ms\\* | 45ms\\* | \\*Includes ~30ms network |\n| Server (Lambda cold start) | 2.8s | 2.8s | Warm-up penalty |\n\nThe M3 Pro numbers are noteworthy. Apple's WebGPU implementation is genuinely excellent, their Metal backend for WebGPU is mature and well-optimized. The RTX 3060 numbers are slightly higher because Chrome's Dawn backend adds a thin Vulkan translation layer.\n\nThe real win isn't the raw latency, it's the elimination of cold starts and network variability. A P99 latency of 25ms in the browser beats a P99 of 3 seconds from a serverless function that might need to cold-start."},{"heading":"Memory Layout Pitfalls","level":2,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#memory-layout-pitfalls","content":"A common pitfall worth highlighting: WGSL's storage buffer alignment rules are strict and poorly documented. Every struct member must be aligned to a multiple of its size, and the struct itself must be aligned to the largest member alignment. For a mixed f32/u32 struct:\n\n```wgsl\nstruct LayerParams {\n weight_offset: u32, // offset 0, align 4\n bias_offset: u32, // offset 4, align 4\n scale: f32, // offset 8, align 4\n num_elements: u32, // offset 12, align 4\n // Total: 16 bytes, align 4, this is fine\n}\n\n// But if you add a vec4:\nstruct LayerParamsV2 {\n weight_offset: u32, // offset 0\n // 12 bytes padding!\n params: vec4<f32>, // offset 16, align 16\n num_elements: u32, // offset 32\n // 12 bytes padding!\n // Total: 48 bytes\n}\n```\n\nThat padding is silent and brutal. Your JavaScript-side buffer writes will be wrong if you don't account for it. Use ` offset_of ` annotations or just avoid vec types in uniform structs."},{"heading":"When NOT to Use This Approach","level":2,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#when-not-to-use-this-approach","content":"There are important limitations to consider. This approach is not suitable when:\n\n- **Your model is larger than ~500M parameters.** GPU memory in browser contexts is limited, and you're sharing it with the renderer.\n- **You need deterministic results across devices.** Floating-point behavior varies between GPU vendors in ways that matter for some applications.\n- **Your users are on integrated graphics.** Intel UHD 630 can run these shaders, but at 10x the latency. Feature-detect and have a fallback.\n- **You need to keep model weights private.** Everything shipped to the browser is inspectable. If your model is proprietary, this isn't for you."},{"heading":"What's Next","level":2,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#whats-next","content":"An active area of exploration is WebGPU implementation of flash attention for longer context windows. The challenge is that WGSL doesn't give you the fine-grained shared memory control that CUDA does, so the tiling strategy needs to be different. This will be worth examining once meaningful benchmarks become available.\n\nThe broader point is this: the industry spent a decade assuming that ML inference belongs on servers. For a growing category of models and use cases, the client GPU is not just adequate, it's superior. Zero latency variance, zero cold starts, zero server costs. The tooling is finally catching up to make this practical."},{"heading":"Frequently asked questions","level":2,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#questions","content":""},{"heading":"Can WebGPU really compete with server-side inference for production ML?","level":3,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#question-1","content":"For models under ~500M parameters with proper quantization, absolutely. We measured 23ms inference on a quantized DistilBERT with WebGPU compute shaders on a mid-range laptop GPU, which is competitive with a cold-start serverless function when you factor in network latency."},{"heading":"What browsers support WebGPU compute shaders right now?","level":3,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#question-2","content":"As of early 2026, Chrome 121+, Edge 121+, and Firefox 130+ ship with full WebGPU compute support. Safari has partial support behind a flag. For production, you should feature-detect with navigator.gpu and fall back to ONNX Runtime Web with WebGL."},{"heading":"How does WGSL compare to GLSL for compute workloads?","level":3,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#question-3","content":"WGSL is more explicit about memory layout and has better safety guarantees. You lose some of the syntactic sugar from GLSL compute shaders, but you gain predictable behavior across vendors. The type system catches buffer layout mismatches at pipeline creation time rather than at runtime."},{"heading":"Tips","level":2,"url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#article-tips","content":"1. **Workgroup sizing matters more than you think:** On most GPUs, ` @workgroup_size(256, 1, 1) ` outperforms ` (64, 4, 1) ` for 1D reductions even though the thread count is the same. The memory coalescing pattern changes drastically with the X dimension.\n2. **Avoid dynamic indexing in WGSL tight loops:** Use ` let cached = weights[idx]; ` outside the inner loop. WGSL doesn't guarantee the compiler will hoist storage buffer reads, and you'll pay the latency on every iteration.\n3. **Double-buffer your storage bindings:** Instead of a single output buffer with a barrier, alternate between ` @group(0) @binding(0) ` and ` @group(0) @binding(1) ` across dispatch passes. This lets the GPU overlap compute and memory operations on architectures that support it.\n4. **Profile with ` GPUCommandEncoder.writeTimestamp() `:** Wrap your dispatch calls with timestamp queries on a dedicated ` GPUQuerySet ` of type ` 'timestamp' `. Resolve into a mappable buffer and read back, this gives you actual GPU-side microsecond timings, not CPU-side noise.\n\n- webgpu\n- wgsl\n- machine-learning\n- browser-inference\n- compute-shaders\n\n[Back to Blog](<https://lucioduran.com/blog>)"}],"links":[{"text":"Español","url":"https://lucioduran.com/es/blog/webgpu-compute-shaders-ml-inference"},{"text":"ONNX Runtime Web","url":"https://onnxruntime.ai/"},{"text":"Back to Blog","url":"https://lucioduran.com/blog"}],"structuredData":[{"@context":"https://schema.org","@graph":[{"@type":"BlogPosting","@id":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference#article","mainEntityOfPage":{"@type":"WebPage","@id":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference"},"headline":"WebGPU Compute Shaders for ML Inference: Client-Side Model Execution","description":"Running ML models directly in the browser with WebGPU compute shaders. WGSL implementation details, benchmarks against WebGL and server-side inference, and practical limitations of client-side model execution.","datePublished":"2025-10-01","author":{"@type":"Person","@id":"https://lucioduran.com/#person","name":"Lucio Durán","url":"https://lucioduran.com"},"publisher":{"@type":"Person","@id":"https://lucioduran.com/#person","name":"Lucio Durán","url":"https://lucioduran.com"},"image":"https://lucioduran.com/favicons/og-default.png","url":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference","inLanguage":"en","isPartOf":{"@id":"https://lucioduran.com/blog#blog"},"articleSection":"WebGPU","keywords":"webgpu, wgsl, machine-learning, browser-inference, compute-shaders","wordCount":1721,"timeRequired":"PT9M","workTranslation":[{"@type":"BlogPosting","inLanguage":"es","url":"https://lucioduran.com/es/blog/webgpu-compute-shaders-ml-inference"}]},{"@type":"BreadcrumbList","itemListElement":[{"@type":"ListItem","position":1,"name":"Home","item":"https://lucioduran.com/"},{"@type":"ListItem","position":2,"name":"Blog","item":"https://lucioduran.com/blog"},{"@type":"ListItem","position":3,"name":"WebGPU Compute Shaders for ML Inference: Client-Side Model Execution","item":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference"}]},{"@type":"FAQPage","mainEntity":[{"@type":"Question","name":"Can WebGPU really compete with server-side inference for production ML?","acceptedAnswer":{"@type":"Answer","text":"For models under ~500M parameters with proper quantization, absolutely. We measured 23ms inference on a quantized DistilBERT with WebGPU compute shaders on a mid-range laptop GPU, which is competitive with a cold-start serverless function when you factor in network latency."}},{"@type":"Question","name":"What browsers support WebGPU compute shaders right now?","acceptedAnswer":{"@type":"Answer","text":"As of early 2026, Chrome 121+, Edge 121+, and Firefox 130+ ship with full WebGPU compute support. Safari has partial support behind a flag. For production, you should feature-detect with navigator.gpu and fall back to ONNX Runtime Web with WebGL."}},{"@type":"Question","name":"How does WGSL compare to GLSL for compute workloads?","acceptedAnswer":{"@type":"Answer","text":"WGSL is more explicit about memory layout and has better safety guarantees. You lose some of the syntactic sugar from GLSL compute shaders, but you gain predictable behavior across vendors. The type system catches buffer layout mismatches at pipeline creation time rather than at runtime."}}]}]}],"contentHash":"75fca63766c96e6331cf9cbc92d575fc547caa3efdc9a2f573db53493cae554d","provenance":{"method":"Deterministic serialization of the server-rendered main element","source":"https://lucioduran.com/blog/webgpu-compute-shaders-ml-inference","hashAlgorithm":"sha256","hashScope":"JSON document fields before contentHash and provenance","sourceHTMLBytes":66127,"contentBytes":13979}}