Skip to content

Feature Request: Add WGSL (WebGPU) Backend Support #37

@TheJoshBrod

Description

@TheJoshBrod

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
        pass

A 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

  • WGSLBackend is registered and selectable in the pipeline
  • A valid .wgsl kernel can be generated end-to-end from a prompt
  • Profiling works via wgpu in Python
  • Validation catches and surfaces syntax errors in generated WGSL

References

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions