-
Notifications
You must be signed in to change notification settings - Fork 0
Feature Request: Add WGSL (WebGPU) Backend Support #37
Description
Add WGSL backend for WebGPU support
The project currently only supports CUDA, with incomplete stubs for Metal and Triton. We should add a proper WGSL backend to support WebGPU, which would let the optimizer run across a much wider range of environments without being tied to NVIDIA hardware.
Background
WebGPU runs on DirectX, Vulkan, and Metal depending on the platform, which means a single WGSL backend covers Windows, macOS, Linux, and mobile. Beyond portability, it opens up the possibility of running kernels directly in the browser for client-side inference in the dashboard. WGSL is also increasingly the right long-term bet for portable GPU compute as WebGL gets phased out.
What needs to change
Create src/optimizer/backends/wgsl/wgsl.py with a WGSLBackend class that implements the Backend interface. It needs three things: a way to query adapter capabilities (max workgroup invocations, etc.) using wgpu, shader module compilation for catching validation errors, and a profiling path that runs a compute shader and measures execution time.
from src.optimizer.core.backend import Backend
class WGSLBackend(Backend):
def get_device_specs(self) -> GPUSpecs:
# Query adapter limits via wgpu (maxComputeInvocationsPerWorkgroup, etc.)
pass
def validate_kernel(self, code: str, ...) -> Tuple[bool, str]:
# Compile shader module via wgpu-py and surface validation errors
pass
def profile_kernel(self, ...) -> dict:
# Run compute shader and measure wall time
passA few related changes are also needed. GPUSpecs in src/optimizer/core/types.py currently assumes warp size as a CUDA concept -- this should either be made optional or mapped to subgroup size for non-CUDA backends. Add src/generator/prompts/wgsl_prompts.py to handle WGSL-specific syntax like @compute, @workgroup_size, and its stricter type system. Register the new backend in pipeline.py. Finally, add wgpu to requirements.txt and make sure a headless adapter works in CI.
Acceptance criteria
WGSLBackendis registered and selectable in the pipeline- A valid
.wgslkernel can be generated end-to-end from a prompt - Profiling works via
wgpuin Python - Validation catches and surfaces syntax errors in generated WGSL
References
- WebGPU spec: https://www.w3.org/TR/webgpu/
- wgpu-py docs: https://wgpu-py.readthedocs.io/