YMRohit's picture
Add 3 recorded verified mints + Codex brief for instant-demo path
59c93fb
Raw
History Blame Contribute Delete
1.69 kB
{"op":"softmax","verified":true,"speedup_compile":1.052,"speedup_eager":4.84,"source":"@triton.jit\ndef _k(x_ptr, y_ptr, stride, N, BLOCK: tl.constexpr):\n row = tl.program_id(0); x_ptr += row * stride; y_ptr += row * stride\n mx = tl.full([BLOCK], -float(\"inf\"), dtype=tl.float32)\n for off in range(0, N, BLOCK):\n cols = off + tl.arange(0, BLOCK)\n x = tl.load(x_ptr + cols, mask=cols < N, other=-float(\"inf\")).to(tl.float32)\n mx = tl.maximum(mx, x)\n rmax = tl.max(mx)\n d = tl.zeros([BLOCK], dtype=tl.float32)\n for off in range(0, N, BLOCK):\n cols = off + tl.arange(0, BLOCK)\n x = tl.load(x_ptr + cols, mask=cols < N, other=-float(\"inf\")).to(tl.float32)\n d += tl.where(cols < N, tl.exp(x - rmax), 0.0)\n den = tl.sum(d)\n for off in range(0, N, BLOCK):\n cols = off + tl.arange(0, BLOCK); m = cols < N\n x = tl.load(x_ptr + cols, mask=m, other=0.0).to(tl.float32)\n tl.store(y_ptr + cols, tl.exp(x - rmax) / den, mask=m)\ndef run(x):\n M, N = x.shape; y = torch.empty_like(x)\n _k[(M,)](x, y, x.stride(0), N, BLOCK=512, num_warps=4)\n return y","statuses":["ok","ok","ok","ok"],"gen_seconds":81.3,"explanation":"Verified correct, and 4.8x faster than plain PyTorch (1.05x vs torch.compile). It even beats torch.compile, the production compiler. A GPU 'kernel' is the tiny program that does one piece of a neural network on the GPU; fusing several steps into one kernel means fewer trips to memory, which is most of what makes AI models faster to run. You just made one — verified by an un-gameable referee, not by anyone's word.","n_verified":4,"k":4,"on_leaderboard":true}