Lab 20 · Browser WASM AI Runtime

Pyodide v0.26.2 WebAssembly Sandbox

WebAssembly AI Runtime & Micrograd Sandbox

Execute genuine AI engineering algorithms entirely client-side inside the Pyodide WebAssembly runtime. Trace forward activations and backpropagate exact gradients across an interactive scalar autograd DAG, verify FlashAttention-3 online softmax stability under extreme logits ($\\pm 1000.0$) with online rescaling $\\alpha$, and translate PagedAttention logical-to-physical block table memory offsets with zero server overhead.

Paper: FlashAttention-3 Guide: GPU Infrastructure System: Local LLM

Pyodide v0.26.2 WASM

Zero-Server-Cost Client Runtime

Browser WebAssembly AI Runtime & Micrograd Sandbox

Challenge:

autograd

beginner

Autograd DAG Enabled

Micrograd Scalar Autograd: Computational Graph & Backpropagation

Micrograd Scalar Autograd Engine

Implement a scalar autograd engine inspired by Andrej Karpathy's Micrograd.

Problem Formulation

A single artificial neuron computes the forward activation:

z = \tanh(w_1 x_1 + w_2 x_2 + b)

Where:

x_1 = 2.0

x_2 = 0.0

(input features)

w_1 = -3.0

w_2 = 1.0

(synaptic weights)

b = 6.8813735870195432

(bias)

Backpropagation Chain Rule

To compute gradients backwards:

\frac{\partial z}{\partial z} = 1.0

\frac{\partial z}{\partial n} = 1 - \tanh^2(n)

\frac{\partial z}{\partial (x_1 w_1)} = \frac{\partial z}{\partial n} \cdot 1, \quad \frac{\partial z}{\partial (x_2 w_2)} = \frac{\partial z}{\partial n} \cdot 1

\frac{\partial z}{\partial x_1} = \frac{\partial z}{\partial (x_1 w_1)} \cdot w_1, \quad \frac{\partial z}{\partial w_1} = \frac{\partial z}{\partial (x_1 w_1)} \cdot x_1

\frac{\partial z}{\partial x_2} = \frac{\partial z}{\partial (x_2 w_2)} \cdot w_2, \quad \frac{\partial z}{\partial w_2} = \frac{\partial z}{\partial (x_2 w_2)} \cdot x_2

Your Task

Complete __mul__ backward propagation: self.grad += other.data * out.grad and other.grad += self.data * out.grad.

Complete tanh backward propagation: self.grad += (1.0 - t**2) * out.grad.

Complete backward() topological sort order traversal to propagate gradients backwards to leaf parameters.

Verify the 8 test assertions and inspect the resulting visual computational DAG.

main.py

Python 3.12 (WASM)

import math

class Value:
    def __init__(self, data, _children=(), _op='', label=''):
        self.data = float(data)
        self.grad = 0.0
        self._backward = lambda: None
        self._prev = set(_children)
        self._op = _op
        self.label = label

def __repr__(self):
        return f"Value(data={self.data:.4f}, grad={self.grad:.4f})"

def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), '+')
        def _backward():
            self.grad += 1.0 * out.grad
            other.grad += 1.0 * out.grad
        out._backward = _backward
        return out

def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), '*')
        # TODO: Implement multiplication backward propagation
        def _backward():
            pass
        out._backward = _backward
        return out

def tanh(self):
        x = self.data
        t = (math.exp(2*x) - 1) / (math.exp(2*x) + 1)
        out = Value(t, (self,), 'tanh')
        # TODO: Implement tanh derivative: (1 - t**2) * out.grad
        def _backward():
            pass
        out._backward = _backward
        return out

def backward(self):
        # TODO: Build topological sort order and call _backward() in reverse order
        pass

# ── Construct Neuron Computational Graph ──
x1 = Value(2.0, label='x1')
x2 = Value(0.0, label='x2')
w1 = Value(-3.0, label='w1')
w2 = Value(1.0, label='w2')
b = Value(6.8813735870195432, label='b')

x1w1 = x1 * w1; x1w1.label = 'x1*w1'
x2w2 = x2 * w2; x2w2.label = 'x2*w2'
x1w1x2w2 = x1w1 + x2w2; x1w1x2w2.label = 'x1*w1 + x2*w2'
n = x1w1x2w2 + b; n.label = 'n'
o = n.tanh(); o.label = 'o'

o.backward()
print(f"Forward output o = {o.data:.4f}, o.grad = {o.grad:.4f}")

Run:Ctrl+Enter

Computation Graph DAG(Auto-extracted from Pyodide execution)

Trace:

10 Nodes·9 Edges

100%

Forward Data

Backward Grad

Test Console & Output

Forward Activationpending

o.data ≈ 0.70710678 (forward pass tanh neuron)

Output Gradientpending

o.grad == 1.0 (seed gradient dz/do)

Pre-Activation Gradientpending

n.grad ≈ 0.5000 (dz/dn = 1 - tanh²(n))

Bias Gradientpending

b.grad ≈ 0.5000 (dz/db = dz/dn × 1)

Input Feature x1 Gradientpending

x1.grad ≈ -1.5000 (dz/dx1 = w1 × dz/dn)

Synaptic Weight w1 Gradientpending

w1.grad ≈ 1.0000 (dz/dw1 = x1 × dz/dn)

Input Feature x2 Gradientpending

x2.grad ≈ 0.5000 (dz/dx2 = w2 × dz/dn)

Synaptic Weight w2 Gradientpending

w2.grad ≈ 0.0000 (dz/dw2 = x2 × dz/dn)

Micrograd Topological Autograd

Construct artificial neuron $z = \\tanh(w_1 x_1 + w_2 x_2 + b)$. Notice how reverse topological traversal recursively triggers _backward(), propagating gradient dz/dw1 = 1.0 and dz/dx1 = -1.5 directly to input weights.

Verifies Topological Backpropagation

Extreme Logits Rescaling (±1000.0)

Standard softmax overflows at +709.7, producing Infinity or NaN. FlashAttention updates running maximum m_curr and applies rescaling factor alpha = exp(m_prev - m_curr) ≤ 1.0 to ensure zero precision loss.

Verifies Online Rescaling Invariant

PagedAttention Memory Translation

Eliminate memory fragmentation by mapping token indices to non-contiguous GPU physical blocks. Translate virtual offset t to block slot phys_block × B_s + (t mod B_s) and enforce runtime page fault handling upon out-of-bounds tokens.

Verifies Virtual-to-Physical Mapping