How to use Vajra
A complete guide to hardware-verified WebGPU kernel optimization with NVIDIA Nemotron.
What is Vajra
“AI writes GPU shaders, but only your hardware knows the truth. Vajra benchmarks every Nemotron kernel live on your GPU and keeps the fastest.”
What is WebGPU / WGSL
WebGPU is the browser API that lets a web page run code on the user's GPU, no install.
WGSL is the shader language for that code.
Vajra tunes WGSL kernels.
The Problem
The same shader is fast on one GPU and slow on another (AMD vs Apple vs NVIDIA). Hand-tuning takes hours and guessing risks wrong maths. AI alone can write code that is slow or wrong — with no check.
Different warp/wavefront widths (32 threads on NVIDIA vs 64 threads on AMD RDNA), register pressure limits, and L1 cache sizes cause identical code to perform wildly differently across chipsets.
Tuning workgroup sizes, unroll factors, vectorization, and shared-memory tile padding requires testing dozens of subtle permutations per hardware target.
LLMs write convincing shader code that introduces data races, missing barriers, out-of-bounds indexing, or subtle precision drift that crashes the pipeline or corrupts neural weights.
How Vajra Works
Five automated steps from prompt to verified, hardware-native speedup:
- 1Select workload: Pick a computationally intensive primitive (SGEMM, Softmax, Conv1D, or Sobel Filter) to optimize.
- 2Run naive baseline: Vajra compiles the reference unoptimized shader and benchmarks it live on your local GPU to establish the ground-truth timing and floating-point reference.
- 3Nemotron writes 3 variants: NVIDIA Nemotron (via Nebius Token Factory or NVIDIA NIM) generates 3 diverse WGSL optimization hypotheses (e.g. shared memory tiling, register caching, loop unrolling).
- 4Compile + verify to 1e-4 + benchmark 9-run median on YOUR GPU: Each generated kernel is compiled directly in your browser. Vajra compares computed output against baseline ground truth (max allowed diff ≤ 1e-4) and measures execution time across 9 GPU passes using high-resolution timestamp queries.
- 5Fastest verified wins, slow/wrong variants go to Failure Graveyard: The fastest numerically verified variant is declared champion with 1-click “Copy Fastest WGSL”. Any code that failed compilation or exceeded 1e-4 tolerance is isolated in the Failure Graveyard.
Modes Explained
Vajra provides four distinct operating modes to accommodate live hackathon submissions, local testing, and non-WebGPU devices:
| Mode | How it works | Benchmarking | Status / Role |
|---|---|---|---|
| Mock | Pre-written variants, benchmarking is live. | Live on your GPU (9-run median) | Default demonstration, instant, no API key required. |
| Live Token Factory | Nemotron generates variants via Nebius Token Factory. | Live on your GPU (9-run median) | Official submission path for hackathon evaluation. |
| Live NVIDIA | Nemotron generates directly via NVIDIA NIM endpoint. | Live on your GPU (9-run median) | Check only, not submission proof. |
| Preview | No WebGPU, recorded reference only. | Pre-recorded reference (Apple M3 Max) | For headless servers or browsers without WebGPU. |
Larger-Project Worked Example: In-Browser LLM / Video Processing
How a production engineering team drops Vajra-tuned kernels into a transformer inference engine or video processing pipeline.
In client-side LLM inference (e.g. running Llama or Whisper in WebGPU) or realtime canvas video processing, matrix multiplication (SGEMM) executes hundreds of times per token or frame.
A 2× to 3× SGEMM speedup compounds directly into reduced time-to-first-token, higher token throughput, and reduced thermal throttling on client laptops and workstations.
@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
let row = global_id.x;
let col = global_id.y;
let N = uniforms.N;
if (row >= N || col >= N) { return; }
var sum: f32 = 0.0;
for (var k: u32 = 0u; k < N; k = k + 1u) {
sum = sum + A[row * N + k] * B[k * N + col];
}
C[row * N + col] = sum;
}var<workgroup> tileA: array<array<f32, 16>, 16>;
var<workgroup> tileB: array<array<f32, 16>, 16>;
@compute @workgroup_size(16, 16)
fn main(
@builtin(local_invocation_id) local_id: vec3<u32>,
@builtin(workgroup_id) group_id: vec3<u32>
) {
let N = uniforms.N;
let row = group_id.x * 16u + local_id.x;
let col = group_id.y * 16u + local_id.y;
var sum: f32 = 0.0;
for (var t: u32 = 0u; t < N; t = t + 16u) {
if (row < N && (t + local_id.y) < N) {
tileA[local_id.x][local_id.y] = A[row * N + (t + local_id.y)];
} else {
tileA[local_id.x][local_id.y] = 0.0;
}
if ((t + local_id.x) < N && col < N) {
tileB[local_id.x][local_id.y] = B[(t + local_id.x) * N + col];
} else {
tileB[local_id.x][local_id.y] = 0.0;
}
workgroupBarrier();
for (var k: u32 = 0u; k < 16u; k = k + 1u) {
sum = sum + tileA[local_id.x][k] * tileB[k][local_id.y];
}
workgroupBarrier();
}
if (row < N && col < N) {
C[row * N + col] = sum;
}
}Exact Paste Location in Your Codebase
Replace the compute kernel entrypoint function in your project file. Keep the exact same @group and @binding buffer harness so your WebGPU pipeline layout, bind groups, and dispatch dimensions remain unchanged. Rebuild your application and run the target inference pipeline.
Note: The above figures are from an individual test run on an AMD RDNA-3 GPU; they are not medians across all devices. Vajra always measures live on your exact silicon.
How to Use Vajra (8 Steps)
Guardrail, in Plain English
“Compile fail, verify fail, and Graveyard mean Vajra refused to ship a wrong result — dropping is a feature, not a bug.”
In machine learning and graphics computation, a shader that runs at 100 GFLOPS but calculates NaN or introduces a 0.01 error will corrupt inference weights, create NaN cascades, and destroy model outputs.
Vajra executes an uncompromising mathematical guardrail: every single variant is compiled on the local WebGPU device and verified against CPU IEEE 754 reference arrays. If the maximum absolute difference exceeds 1e-4, or if the compiler reports a syntax or resource error, the variant is instantly quarantined in the Failure Graveyard.
Under no circumstances will Vajra promote a broken or inaccurate kernel to winner, regardless of how fast it executes.