Vajra: WebGPU Kernel Auto-Tuner

AI writes GPU shaders, but only your silicon knows the truth. Vajra generates WGSL via Nemotron on Nebius, verifies <1e-4 float accuracy, and benchmarks live on your GPU for up to 7.9x speedups.

Comment 1

Inspiration

Language models can write syntactically plausible GPU compute shaders, but an LLM cannot know how that code executes on physical silicon. A model's idea of an "optimized" kernel is a static guess — it has never seen your GPU's cache hierarchy, warp/wave scheduler, memory bus, or driver compiler backend. In practice, manual WGSL tuning means hours of trial-and-error edits in browser DevTools.

We built Vajra (Sanskrit for thunderbolt) to close the loop between cloud LLM code generation and local hardware truth: generate kernel candidates with AI, then compile, verify, and benchmark them directly on the machine in front of you to crown the fastest verified kernel.

What it does

Vajra is a zero-install, browser-native WebGPU kernel auto-tuner powered by NVIDIA Nemotron via Nebius Token Factory:

  • Pick a Compute Workload: Choose from 4 standard compute patterns: Matrix Multiplication (SGEMM 512x512), Row-wise Softmax (1024x1024), 1D Convolution (N = 1,048,576, K = 64), or 2D Sobel Edge Filter (2048x2048).
  • Run the Naive Baseline: Vajra dispatches a deterministic naive baseline shader on your local GPU, capturing reference output buffers and baseline execution latency.
  • Generate 3 Architectural Variants: Nemotron synthesizes three distinct WGSL compute strategies: Workgroup Geometry Tuning, Shared Workgroup Memory Tiling (var), and Vectorized SIMD Math (vec4 / fma).
  • Compile & Guardrail: Each variant is compiled via browser WebGPU (createShaderModule + getCompilationInfo). Broken variants are quarantined in an interactive Failure Graveyard with exact compiler line/column diagnostics without crashing the pipeline.
  • Numerical Verification: Every compiled kernel runs against identical input buffers and is checked against the baseline using strict relative tolerance (max relative error under 1e-4). Inaccurate shaders are rejected.
  • Median Benchmarking & Report Card: Verified kernels undergo 5 warmup dispatches followed by 9 timed dispatches. Vajra reports 9-run medians (never cherry-picked best runs), live GFLOPS bar charts, and a Silicon Report Card detailing speedup vs baseline, derived memory bandwidth (GB/s), workgroup configurations, arithmetic intensity roofline regimes, and timing stability (IQR).
  • Agentic Gen 2 Evolution: A single click evolves the winning kernel into a second-generation candidate (Variant D), feeding measured GFLOPS and Failure Graveyard compiler logs back to Nemotron.
  • Drop-in Export: Developers can inspect the stitched WGSL and copy the verified fastest shader or an exported TypeScript dispatch module in one click.

How we built it

  • Frontend & Systems Architecture: Built with Next.js (App Router) and TypeScript, styled using a high-contrast, flat Vercel-inspired dashboard palette (zero CSS gradients, zero purple or indigo).
  • Deterministic WebGPU Harness: A hardcoded, browser-native execution engine managing deterministic Float32Array seeding, uniform buffer layouts, bind groups, pipeline staging, and readbacks.
  • Inference Chain (Nebius Token Factory): Server-side routing calls NVIDIA Nemotron (Super primary, with automatic failover to Ultra) via OpenAI-compatible endpoints on Nebius Token Factory. API keys live exclusively in server environment variables.
  • Tavily Web Intelligence Integration: Dual-stage runtime grounding queries the Tavily Search API before generation to retrieve hardware-aware W3C WGSL specification rules, while dynamically searching exact Dawn/wgpu compiler errors from the Failure Graveyard to guide Gen 2 error recovery.
  • Dual Timing Engine: Uses native WebGPU hardware timestamp queries (GPUQuerySet 64-bit nanoseconds) where supported, falling back cleanly to wall-clock timing via performance.now() wrapped around queue completion.
  • Honest Hardware Telemetry: Reports only parsed workgroup configurations (@workgroup_size and dispatch grid dimensions) and derived memory throughput (total bytes moved divided by median execution time), never fabricating hardware counters WebGPU cannot expose.

Challenges we ran into

  • WGSL Compiler Incompatibilities & Uniformity Rules: Small models frequently hallucinate non-uniform barriers (such as calling workgroupBarrier after an early return guard). We built an AST/sanitizer pass to hoist boundary checks past barriers, enforced uniform barrier prompt patterns, and designed the Failure Graveyard to cleanly drop illegal shaders without stalling the GPU command queue.
  • Vectorized Dispatch Truncation: When vectorizing row-wise kernels (like Softmax with vec4), models often declare elementsPerThread as [4, 1, 1] for column loops, inadvertently causing the harness to under-dispatch workgroups along X and leaving the tail 75% of rows uncomputed. We clamped outer problem-dimension divisors to ensure 100% of rows dispatch while vectorizing inner loops.
  • Fast-Math Reordering Drift: Different GPU architectures and loop-unrolling strategies accumulate IEEE-754 floating-point values in different summation orders. Exact equality checks falsely fail valid optimized kernels. We implemented an element-wise relative tolerance validator (checking if relative error is under 1e-4), catching mathematical divergence while respecting hardware reordering.
  • Heterogeneous Environments: Some browser sessions lack WebGPU support or run in headless sandboxes. We engineered a seamless Recorded Reference Fallback (benchmarked on Apple Silicon) with full stage simulation and fault-injection testing, ensuring every judge and visitor experiences the complete pipeline regardless of their device.

Accomplishments that we're proud of

  • A Truly Closed Agentic Loop: Built an end-to-end autonomous coding workflow where an AI generates code, physical hardware tests it, compiler errors inform the next generation, and only numerically verified winners are crowned.
  • Measured, Repeatable Hardware Speedups: In live physical testing on an integrated AMD Radeon 780M (RDNA-3) GPU, Vajra consistently unlocked:
  • Matrix Multiplication (SGEMM 512x512): Baseline 30.6 GFLOPS to Variant B 132.9 GFLOPS (4.34x speedup, verified under 1e-4).
  • Row-wise Softmax (1024x1024): Baseline 0.3 GFLOPS to Variant B 2.2 GFLOPS (7.90x speedup, verified under 1e-4).
  • 1D Convolution (N = 1M, K = 64): Baseline 31.5 GFLOPS to Variant A 52.1 GFLOPS (1.65x speedup, verified under 1e-4).

  • 2D Sobel Filter (2048x2048): Baseline 33.7 GFLOPS to Variant A 65.1 GFLOPS (1.93x speedup, verified under 1e-4).

  • Zero-Crash Resilience: Achieved 100% pipeline stability across invalid syntax, memory boundary overflows, and unverified float outputs with zero unhandled WebGPU promise rejections.

What we learned

  • Physical Measurement Trumps Text Generation: An LLM can write plausible-looking code, but a shader that has not been compiled and timed on your exact GPU architecture is merely an unverified hypothesis.
  • Telemetry Honesty Builds Trust: WebGPU does not expose hardware SM occupancy or DRAM bus counters. Clearly labeling metrics as derived and reporting 9-run medians with interquartile stability ranges (IQR) provides actionable engineering value over vanity numbers.
  • Restricted Scoping Empowers Agentic Coding: Constraining the LLM to write strictly the inner compute body — while letting a deterministic TypeScript harness manage memory bindings and synchronization — turns erratic generation into a robust engineering tool.
  • Failure Transparency is a First-Class Feature: Surfacing compiler line errors and numerical drift in the Failure Graveyard proves why hardware-in-the-loop verification is essential.

What's next for Vajra

  • Broader Kernel Workloads: Expand from foundational linear algebra and image stencils to FlashAttention-2 approximations, LayerNorm, and 1D/2D Fast Fourier Transforms (FFT).
  • Subgroup Intrinsics: Leverage emerging WebGPU subgroup operations (subgroupBallot, subgroupShuffle) for warp-level reductions without workgroup shared-memory overhead.
  • Multi-Device Comparison Matrix: Allow developers to export hardware lockfiles (vajra.lock.json) to benchmark and compare identical WGSL shaders across NVIDIA RTX, Apple M-Series, AMD RDNA, and Intel Arc hardware.
  • IDE & Toolchain Integration: Package the Vajra verification harness into a CLI and VS Code extension to auto-tune WGSL and ONNX Runtime Web shaders inline during local build workflows.

Built With

Share this project:

Updates

Submission history