KernelBench cuda · B200
Grid + MinGRU SPS Claude Opus 4.8
85.9%geomean peak fraction across shapes
manually audited: clean
Genuine fused CUDA megastep kernel: one rollout_step launch per env step performs env transition (LCG food respawn folded into the head of the next launch), a layer-0 encoder algebraically folded into the gate projection (768x256 GEMM -> 768x4), two full 768x256 GEMMs on fp16 tensor cores via hand-written mma.sync PTX with an fp32-accurate hi/lo split (3 mma terms), the MinGRU cell + highway, the action head, argmax, movement, reward, and a ballot/atomicOr anyhit reduction. policy_forward and env_step are separate exact-fp32 CUDA kernels matching the graders' component checks. Weight prep, an mt19937 host reimplementation reproducing torch's CPU randint stream bit-for-bit (legitimate re-derivation of the reference's seeded init, not stack sniffing), and a fused init kernel round out the timed path. No output memoization, no input-identity dispatch, no CUDA graphs, no grader/tolerance edits, no forbidden frameworks. Official benchmark: per-shape fractions 0.7173 / 0.8922 / 0.9755 / 0.8756, geomean 0.8599 of the deck's fixed 150M SPS anchor (~129M geomean steps/s on B200).
harnessclaudeagent session2h 13mtotal wall2h 13mcheck29sbenchmark2soutput tokens458,396cost$43.13gpu-lock wait6sgpu-lock held2mregimethroughput
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
No per-shape benchmark data archived for this run.
Kernel source (redacted)
"""Fused CUDA grid-foraging env + 3x MinGRU(h=256) policy rollout.
The rollout runs as one megastep kernel per env step (env transition + encoder +
3 MinGRU layers + action head, all fused). Design notes are in _CUDA_SRC below;
the short version:
* the encoder is algebraically folded into layer 0, turning one of the three
768x256 GEMMs into a 768x4 one (1.49x fewer FLOPs);
* the two remaining GEMMs run on fp16 tensor cores using an fp32-accurate
hi/lo split (3 mma terms), so logits stay within ~3e-8 of the reference;
* 16 warps x 2 resident blocks per SM lets one block's MinGRU cell overlap the
other's GEMM, which is what actually keeps the tensor pipe fed.
policy_forward/env_step are separate exact-fp32 kernels (they are correctness
surfaces, not throughput ones).
"""
from __future__ import annotations
import os
from pathlib import Path
import torch
import torch.nn as nn
BOARD = 11
OBS_DIM = 4
HIDDEN = 256
GRU_LAYERS = 3
NUM_ACTIONS = 4
GRU_OUT = 3 * HIDDEN
_CUDA_SRC = r"""
// Grid-foraging env + 3x MinGRU(h=256) policy rollout -- fused CUDA megastep kernel.
//
// Key ideas
// 1. Layer-0 fold: gates_0 = W0 @ (Wenc @ obs + benc) = (W0 Wenc) @ obs + W0 benc.
// Since obs is 4-dim this turns a 768x256 GEMM into a 768x4 one -> 1.49x fewer FLOPs.
// 2. The two remaining 768x256 GEMMs run on fp16 tensor cores with a hi/lo split
// (3 mma terms) which reproduces fp32 accuracy to ~2e-7 relative.
// 3. Weights are pre-permuted into mma B-fragment order so the inner loop is one
// perfectly coalesced 128-bit global load per (k-step, n-tile).
// 4. One kernel launch per env step; the previous step's food respawn is folded into
// the head of the next launch so the global `hit.any()` reduction costs nothing
// beyond the kernel boundary that is already there.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <vector>
#include <cstring>
#if defined(__x86_64__)
#include <immintrin.h>
#endif
#define BOARD 11
#define HID 256
#define NGATE 768
#define NACT 4
#define LCG_A 6364136223846793005ULL
#define LCG_MASK 0x7FFFFFFFFFFFFFFFLL
// hi/lo split scales: keeps the fp16 "lo" limbs inside the normal range.
#define ASCALE 1024.0f
#define BSCALE 1024.0f
#define INVSCALE (1.0f / (ASCALE * BSCALE))
#define RENV 32 // envs per block
#define NWARP 16 // warps per block (2 blocks/SM at 512 threads, 64 regs)
#define NTHREADS (NWARP * 32)
#define HSTRIDE 264 // halves per row of the shared h_hi / h_lo staging arrays
#define GPW 1 // 8-wide hidden groups handled per warp per pass
#define PASSES (32 / (NWARP * GPW))
#define NTPW (3 * GPW) // n-tiles (of 8 gate columns) handled per warp
#define MTILE (RENV / 16)
#define BPSM 2 // resident blocks per SM
// --------------------------------------------------------------------------------------
// PTX helpers
// --------------------------------------------------------------------------------------
__device__ __forceinline__ void mma_m16n8k16(float *d, const uint32_t *a, const uint32_t *b) {
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
"{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
: "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
}
__device__ __forceinline__ void ldmatrix_x4(uint32_t *r, const __half *p) {
uint32_t addr = static_cast<uint32_t>(__cvta_generic_to_shared(p));
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
: "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
: "r"(addr));
}
__device__ __forceinline__ float sigmoidf(float x) { return 1.0f / (1.0f + expf(-x)); }
// Throughput variants for the rollout: 4 instructions per gate instead of ~12/~25.
// The caller folds the accumulator rescale into the exponent constant, so
// sigma(s*z) is rcp(1 + ex2(z * c)) with c = -s*log2(e). Absolute error is
// ~5e-8 (sigmoid) / ~1.2e-7 (tanh) -- far below the 1e-3 logit tolerance.
#define LOG2E 1.4426950408889634f
// .ftz is safe here: the exponent argument never lands in the subnormal range
// (|z*c| stays well inside +-150) and rcp's operand is always >= 1.
__device__ __forceinline__ float ex2a(float x) {
float r;
asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x));
return r;
}
__device__ __forceinline__ float rcpa(float x) {
float r;
asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x));
return r;
}
__device__ __forceinline__ float sigc(float z, float c) { return rcpa(1.0f + ex2a(z * c)); }
__device__ __forceinline__ float tanhc(float z, float c) {
return fmaf(2.0f, sigc(z, 2.0f * c), -1.0f);
}
__device__ __forceinline__ uint32_t pack2(__half a, __half b) {
__half2 h = __halves2half2(a, b);
return *reinterpret_cast<uint32_t *>(&h);
}
// --------------------------------------------------------------------------------------
// Weight preprocessing: (768,256) fp32 -> mma B-fragment order, fp16 hi/lo.
//
// Permuted gate column c maps to original row: c = 24*gh + 8*gate + jj -> gate*256+gh*8+jj
// so that the three gates of one 8-wide hidden group are three consecutive n-tiles.
// Output index: ((kstep*96 + ntile)*32 + lane), one uint4 = {hi01, hi89, lo01, lo89}.
// --------------------------------------------------------------------------------------
__global__ void permute_weights(const float *__restrict__ W, uint4 *__restrict__ out,
int total) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total) return;
const int layer = idx / (16 * 96 * 32);
const int sub = idx - layer * (16 * 96 * 32);
W += (size_t)layer * NGATE * HID;
int lane = sub & 31;
int ntile = (sub >> 5) % 96;
int kstep = sub / (96 * 32);
int c = ntile * 8 + (lane >> 2);
int gh = c / 24, rem = c - gh * 24, gate = rem >> 3, jj = rem & 7;
int grow = gate * HID + gh * 8 + jj;
int k = kstep * 16 + ((lane & 3) << 1);
const float *wp = W + (size_t)grow * HID;
float v[4] = {wp[k] * BSCALE, wp[k + 1] * BSCALE, wp[k + 8] * BSCALE, wp[k + 9] * BSCALE};
__half hi[4], lo[4];
#pragma unroll
for (int i = 0; i < 4; i++) {
hi[i] = __float2half_rn(v[i]);
lo[i] = __float2half_rn(v[i] - __half2float(hi[i]));
}
uint4 o;
o.x = pack2(hi[0], hi[1]);
o.y = pack2(hi[2], hi[3]);
o.z = pack2(lo[0], lo[1]);
o.w = pack2(lo[2], lo[3]);
out[idx] = o;
}
// Build the permuted (4,768) / (768,) layer-0 folded constants from natural order.
// Fold the encoder into layer 0's gate weights: M0 = W0 @ Wenc, c0 = W0 @ benc.
// fp64 accumulation keeps this below the fp32 round-off of the reference chain.
__global__ __launch_bounds__(256) void fold_l0(const float *__restrict__ wg0, // (768,256)
const float *__restrict__ wenc, // (256,4)
const float *__restrict__ benc, // (256,)
float *__restrict__ M0, // (768,4)
float *__restrict__ c0) { // (768,)
__shared__ double red[8][5];
const int r = blockIdx.x, k = threadIdx.x, lane = k & 31, w = k >> 5;
const double wv = (double)wg0[(size_t)r * HID + k];
double v[5];
#pragma unroll
for (int o = 0; o < 4; o++) v[o] = wv * (double)wenc[k * 4 + o];
v[4] = wv * (double)benc[k];
#pragma unroll
for (int off = 16; off; off >>= 1)
#pragma unroll
for (int o = 0; o < 5; o++) v[o] += __shfl_down_sync(0xffffffff, v[o], off);
if (lane == 0)
#pragma unroll
for (int o = 0; o < 5; o++) red[w][o] = v[o];
__syncthreads();
if (k < 5) {
double s = 0;
#pragma unroll
for (int i = 0; i < 8; i++) s += red[i][k];
if (k < 4) M0[r * 4 + k] = (float)s;
else c0[r] = (float)s;
}
}
// Layer-0 fragment table.
// n-tile 0..95 : folded gate weights gates0 = M0 @ obs + c0 (permuted columns)
// n-tile 96..127: encoder weights h0 = Wenc @ obs + benc
// K is padded to 16: rows 0..3 hold the obs weights, row 4 holds the bias (the A
// operand carries a constant 1 there), rows 5..15 are zero.
__global__ void permute_l0(const float *__restrict__ M0, // (768,4)
const float *__restrict__ c0, // (768,)
const float *__restrict__ wenc, // (256,4)
const float *__restrict__ benc, // (256,)
uint4 *__restrict__ out) { // [128][32]
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= 128 * 32) return;
int lane = idx & 31;
int ntile = idx >> 5;
int k0 = (lane & 3) << 1;
const float *wrow;
float bias;
if (ntile < 96) {
int c = ntile * 8 + (lane >> 2);
int gh = c / 24, rem = c - gh * 24, gate = rem >> 3, jj = rem & 7;
int grow = gate * HID + gh * 8 + jj;
wrow = M0 + (size_t)grow * 4;
bias = c0[grow];
} else {
int j = (ntile - 96) * 8 + (lane >> 2);
wrow = wenc + (size_t)j * 4;
bias = benc[j];
}
// k values covered by this lane: k0, k0+1, k0+8, k0+9
float v[4];
const int ks[4] = {k0, k0 + 1, k0 + 8, k0 + 9};
#pragma unroll
for (int i = 0; i < 4; i++) {
int k = ks[i];
v[i] = (k < 4) ? wrow[k] * BSCALE : ((k == 4) ? bias * BSCALE : 0.0f);
}
__half hi[4], lo[4];
#pragma unroll
for (int i = 0; i < 4; i++) {
hi[i] = __float2half_rn(v[i]);
lo[i] = __float2half_rn(v[i] - __half2float(hi[i]));
}
uint4 o;
o.x = pack2(hi[0], hi[1]);
o.y = pack2(hi[2], hi[3]);
o.z = pack2(lo[0], lo[1]);
o.w = pack2(lo[2], lo[3]);
out[idx] = o;
}
// --------------------------------------------------------------------------------------
// Fused rollout step
// --------------------------------------------------------------------------------------
#define NT0 (NTPW + GPW) // layer-0 n-tiles per warp: NTPW gate tiles + GPW encoder tiles
struct SmemPtrs {
__half *hhi; // [2][RENV*HSTRIDE] ping-pong
__half *hlo; // [2][RENV*HSTRIDE]
float *hf; // [RENV*HID] fragment-major fp32 h
float *lred; // [NACT*NWARP*RENV]
uint32_t *obsf; // [RENV][4]: packed layer-0 A fragments (hi,lo) x (lane group 0,1)
};
__device__ __forceinline__ SmemPtrs carve(char *base) {
SmemPtrs s;
char *p = base;
s.hhi = (__half *)p; p += 2 * RENV * HSTRIDE * sizeof(__half);
s.hlo = (__half *)p; p += 2 * RENV * HSTRIDE * sizeof(__half);
s.hf = (float *)p; p += RENV * HID * sizeof(float);
s.lred = (float *)p; p += NACT * NWARP * RENV * sizeof(float);
s.obsf = (uint32_t *)p; p += RENV * 4 * sizeof(uint32_t);
return s;
}
static size_t smem_bytes() {
return 4 * RENV * HSTRIDE * sizeof(__half) + RENV * HID * sizeof(float) +
NACT * NWARP * RENV * sizeof(float) + RENV * 4 * sizeof(uint32_t);
}
// LCG used by the reference env for food respawns.
__device__ __forceinline__ long long lcg(long long r) {
return (long long)((unsigned long long)r * LCG_A + 1ULL) & LCG_MASK;
}
__global__ __launch_bounds__(NTHREADS, BPSM) void rollout_step(
const uint4 *__restrict__ wfrag, // [2][16][96][32]
const uint4 *__restrict__ wfrag0, // [128][32]
const float *__restrict__ wa_g, // (4,256)
const float *__restrict__ ba_g, // (4,)
float *__restrict__ state, // (N,3,256)
int *__restrict__ agent, // (N,2)
int *__restrict__ food, // (N,2)
long long *__restrict__ rng, // (N,)
float *__restrict__ rewards, // (N,)
float *__restrict__ last_logits, // (N,4)
int *__restrict__ anyhit, // (horizon,)
int N, int step) {
extern __shared__ char smem_raw[];
SmemPtrs sm = carve(smem_raw);
const int tid = threadIdx.x;
const int warp = tid >> 5;
const int lane = tid & 31;
const int e0 = blockIdx.x * RENV;
const int prev_any = (step > 0) ? anyhit[step - 1] : 0;
// ---------------------------------------------------------------------------
// Layer-0 A operand: [obs0 obs1 obs2 obs3 1 0 ... 0] * ASCALE, built straight
// into mma fragment registers -- lane group (lane&3) selects which two of the
// 16 padded columns this thread holds, so lanes 2,3 need no env data at all.
// ---------------------------------------------------------------------------
if (tid < RENV) {
const int ge = e0 + tid;
const int valid = (ge < N);
int ax = 0, ay = 0, fx = 1, fy = 1;
if (valid) {
ax = agent[ge * 2 + 0];
ay = agent[ge * 2 + 1];
fx = food[ge * 2 + 0];
fy = food[ge * 2 + 1];
if (prev_any && ax == fx && ay == fy) {
long long r = lcg(rng[ge]);
fx = (int)(r % BOARD);
r = lcg(r);
fy = (int)(r % BOARD);
}
}
const float o[4] = {(float)(fx - ax) / (float)BOARD * ASCALE,
(float)(fy - ay) / (float)BOARD * ASCALE,
(float)ax / (float)(BOARD - 1) * ASCALE,
(float)ay / (float)(BOARD - 1) * ASCALE};
#pragma unroll
for (int g = 0; g < 2; g++) {
__half h0 = __float2half_rn(o[2 * g]), h1 = __float2half_rn(o[2 * g + 1]);
sm.obsf[tid * 4 + g * 2 + 0] = pack2(h0, h1);
sm.obsf[tid * 4 + g * 2 + 1] = pack2(__float2half_rn(o[2 * g] - __half2float(h0)),
__float2half_rn(o[2 * g + 1] - __half2float(h1)));
}
}
__syncthreads();
const int cgrp = lane & 3;
const uint32_t abias = pack2(__float2half_rn(ASCALE), __float2half_rn(0.0f));
uint32_t a0hi[MTILE][4], a0lo[MTILE][4];
#pragma unroll
for (int mt = 0; mt < MTILE; mt++) {
#pragma unroll
for (int half = 0; half < 2; half++) {
uint32_t hi = 0, lo = 0;
if (cgrp < 2) {
const int e = mt * 16 + (lane >> 2) + half * 8;
hi = sm.obsf[e * 4 + cgrp * 2 + 0];
lo = sm.obsf[e * 4 + cgrp * 2 + 1];
} else if (cgrp == 2) {
hi = abias; // constant column feeding the folded bias row
}
a0hi[mt][half] = hi;
a0lo[mt][half] = lo;
}
a0hi[mt][2] = 0; a0hi[mt][3] = 0;
a0lo[mt][2] = 0; a0lo[mt][3] = 0;
}
// ---------------------------------------------------------------------------
// Three MinGRU layers. L=0 is the folded encoder+gate rank-4 product (one
// k-step, NT0 n-tiles); L=1,2 are the full 768x256 GEMMs (16 k-steps).
// ---------------------------------------------------------------------------
const int arow = ((lane >> 3) & 1) * 8 + (lane & 7);
const int acol = ((lane >> 4) & 1) * 8;
const float CE = -LOG2E * INVSCALE;
#pragma unroll
for (int L = 0; L < 3; L++) {
#pragma unroll
for (int ps = 0; ps < PASSES; ps++) {
const int gbase = (ps * NWARP + warp) * GPW;
float acc[NT0][MTILE][4];
#pragma unroll
for (int n = 0; n < NT0; n++)
#pragma unroll
for (int m = 0; m < MTILE; m++)
#pragma unroll
for (int k = 0; k < 4; k++) acc[n][m][k] = 0.0f;
if (L == 0) {
const uint4 *wp = wfrag0 + (size_t)(gbase * 3) * 32 + lane;
#pragma unroll
for (int nt = 0; nt < NTPW; nt++) {
uint4 b = wp[nt * 32];
uint32_t bhi[2] = {b.x, b.y}, blo[2] = {b.z, b.w};
#pragma unroll
for (int mt = 0; mt < MTILE; mt++) {
mma_m16n8k16(acc[nt][mt], a0hi[mt], bhi);
mma_m16n8k16(acc[nt][mt], a0lo[mt], bhi);
mma_m16n8k16(acc[nt][mt], a0hi[mt], blo);
}
}
const uint4 *we = wfrag0 + (size_t)(96 + gbase) * 32 + lane;
#pragma unroll
for (int gg = 0; gg < GPW; gg++) {
uint4 b = we[gg * 32];
uint32_t bhi[2] = {b.x, b.y}, blo[2] = {b.z, b.w};
#pragma unroll
for (int mt = 0; mt < MTILE; mt++) {
mma_m16n8k16(acc[NTPW + gg][mt], a0hi[mt], bhi);
mma_m16n8k16(acc[NTPW + gg][mt], a0lo[mt], bhi);
mma_m16n8k16(acc[NTPW + gg][mt], a0hi[mt], blo);
}
}
} else {
const uint4 *wl = wfrag + (size_t)(L - 1) * (16 * 96 * 32);
const __half *hhi_in = sm.hhi + (size_t)(L - 1) * (RENV * HSTRIDE);
const __half *hlo_in = sm.hlo + (size_t)(L - 1) * (RENV * HSTRIDE);
const uint4 *wp0 = wl + (size_t)(gbase * 3) * 32 + lane;
// software pipeline: the B fragment for k-step ks+1 is issued as soon as
// the one for ks has been consumed.
uint4 bcur[NTPW];
#pragma unroll
for (int nt = 0; nt < NTPW; nt++) bcur[nt] = wp0[nt * 32];
for (int ks = 0; ks < 16; ks++) {
uint32_t ahi[MTILE][4], alo[MTILE][4];
#pragma unroll
for (int mt = 0; mt < MTILE; mt++) {
ldmatrix_x4(ahi[mt], hhi_in + (mt * 16 + arow) * HSTRIDE + ks * 16 + acol);
ldmatrix_x4(alo[mt], hlo_in + (mt * 16 + arow) * HSTRIDE + ks * 16 + acol);
}
const uint4 *wn = wp0 + (size_t)(ks < 15 ? ks + 1 : 15) * (96 * 32);
uint4 bv[NTPW];
#pragma unroll
for (int nt = 0; nt < NTPW; nt++) {
bv[nt] = bcur[nt];
bcur[nt] = wn[nt * 32];
}
// Issue the three emulation terms as three separate sweeps so that
// successive mma on the same accumulator are NTPW*MTILE apart.
#pragma unroll
for (int nt = 0; nt < NTPW; nt++) {
uint32_t bhi[2] = {bv[nt].x, bv[nt].y};
#pragma unroll
for (int mt = 0; mt < MTILE; mt++) mma_m16n8k16(acc[nt][mt], ahi[mt], bhi);
}
#pragma unroll
for (int nt = 0; nt < NTPW; nt++) {
uint32_t bhi[2] = {bv[nt].x, bv[nt].y};
#pragma unroll
for (int mt = 0; mt < MTILE; mt++) mma_m16n8k16(acc[nt][mt], alo[mt], bhi);
}
#pragma unroll
for (int nt = 0; nt < NTPW; nt++) {
uint32_t blo[2] = {bv[nt].z, bv[nt].w};
#pragma unroll
for (int mt = 0; mt < MTILE; mt++) mma_m16n8k16(acc[nt][mt], ahi[mt], blo);
}
}
}
// ---- MinGRU cell + highway ----
__half *hhi_out = sm.hhi + (size_t)(L & 1) * (RENV * HSTRIDE);
__half *hlo_out = sm.hlo + (size_t)(L & 1) * (RENV * HSTRIDE);
#pragma unroll
for (int gg = 0; gg < GPW; gg++) {
const int jc = (gbase + gg) * 8 + ((lane & 3) << 1);
#pragma unroll
for (int mt = 0; mt < MTILE; mt++) {
#pragma unroll
for (int half = 0; half < 2; half++) {
const int e = mt * 16 + (lane >> 2) + half * 8;
const int ge = e0 + e;
const int s0 = half * 2;
const int hidx = ((gbase + gg) * MTILE + mt) * 128 + s0 * 32 + lane;
float th0 = tanhc(acc[gg * 3 + 0][mt][s0], CE);
float th1 = tanhc(acc[gg * 3 + 0][mt][s0 + 1], CE);
float sg0 = sigc(acc[gg * 3 + 1][mt][s0], CE);
float sg1 = sigc(acc[gg * 3 + 1][mt][s0 + 1], CE);
float p0 = sigc(acc[gg * 3 + 2][mt][s0], CE);
float p1 = sigc(acc[gg * 3 + 2][mt][s0 + 1], CE);
float hp0, hp1;
if (L == 0) {
hp0 = acc[NTPW + gg][mt][s0] * INVSCALE;
hp1 = acc[NTPW + gg][mt][s0 + 1] * INVSCALE;
} else {
hp0 = sm.hf[hidx];
hp1 = sm.hf[hidx + 32];
}
float *sp = state + ((size_t)ge * 3 + L) * HID + jc;
float2 st2 = make_float2(0.f, 0.f);
if (ge < N) st2 = *reinterpret_cast<float2 *>(sp);
float o0 = st2.x + sg0 * (th0 - st2.x);
float o1 = st2.y + sg1 * (th1 - st2.y);
float hn0 = fmaf(p0, o0 - hp0, hp0);
float hn1 = fmaf(p1, o1 - hp1, hp1);
if (ge < N) *reinterpret_cast<float2 *>(sp) = make_float2(o0, o1);
sm.hf[hidx] = hn0;
sm.hf[hidx + 32] = hn1;
if (L < 2) {
__half2 hi =
__halves2half2(__float2half_rn(hn0 * ASCALE), __float2half_rn(hn1 * ASCALE));
__half2 lo = __halves2half2(
__float2half_rn(hn0 * ASCALE - __half2float(__low2half(hi))),
__float2half_rn(hn1 * ASCALE - __half2float(__high2half(hi))));
*reinterpret_cast<__half2 *>(hhi_out + e * HSTRIDE + jc) = hi;
*reinterpret_cast<__half2 *>(hlo_out + e * HSTRIDE + jc) = lo;
}
}
}
}
}
if (L < 2) __syncthreads();
}
// ---- logits = wa h3 + ba ----
{
float part[MTILE][2][NACT];
#pragma unroll
for (int mt = 0; mt < MTILE; mt++)
#pragma unroll
for (int half = 0; half < 2; half++)
#pragma unroll
for (int a = 0; a < NACT; a++) part[mt][half][a] = 0.0f;
#pragma unroll
for (int pg = 0; pg < PASSES * GPW; pg++) {
const int grp = (pg / GPW) * NWARP * GPW + warp * GPW + (pg % GPW);
const int jc = grp * 8 + ((lane & 3) << 1);
float wa0[NACT], wa1[NACT];
#pragma unroll
for (int a = 0; a < NACT; a++) {
wa0[a] = wa_g[a * HID + jc];
wa1[a] = wa_g[a * HID + jc + 1];
}
#pragma unroll
for (int mt = 0; mt < MTILE; mt++)
#pragma unroll
for (int half = 0; half < 2; half++) {
const int hidx = (grp * MTILE + mt) * 128 + half * 64 + lane;
float h0 = sm.hf[hidx];
float h1 = sm.hf[hidx + 32];
#pragma unroll
for (int a = 0; a < NACT; a++) part[mt][half][a] += wa0[a] * h0 + wa1[a] * h1;
}
}
#pragma unroll
for (int mt = 0; mt < MTILE; mt++)
#pragma unroll
for (int half = 0; half < 2; half++)
#pragma unroll
for (int a = 0; a < NACT; a++) {
float v = part[mt][half][a];
v += __shfl_xor_sync(0xffffffff, v, 1);
v += __shfl_xor_sync(0xffffffff, v, 2);
if ((lane & 3) == 0) {
const int e = mt * 16 + (lane >> 2) + half * 8;
sm.lred[(a * NWARP + warp) * RENV + e] = v;
}
}
}
__syncthreads();
// ---- action, env transition, reward ----
if (tid < RENV) {
const int e = tid;
const int ge = e0 + e;
const int valid = (ge < N);
float lg[NACT];
#pragma unroll
for (int a = 0; a < NACT; a++) {
float v = ba_g[a];
for (int w = 0; w < NWARP; w++) v += sm.lred[(a * NWARP + w) * RENV + e];
lg[a] = v;
}
int act = 0;
float best = lg[0];
#pragma unroll
for (int a = 1; a < NACT; a++)
if (lg[a] > best) { best = lg[a]; act = a; }
int ax = 0, ay = 0, fx = 1, fy = 1;
long long r = 0;
if (valid) {
ax = agent[ge * 2 + 0];
ay = agent[ge * 2 + 1];
fx = food[ge * 2 + 0];
fy = food[ge * 2 + 1];
r = rng[ge];
if (prev_any) {
long long r1 = lcg(r);
long long r2 = lcg(r1);
if (ax == fx && ay == fy) {
fx = (int)(r1 % BOARD);
fy = (int)(r2 % BOARD);
}
r = r2;
}
}
if (act == 0) ay -= 1;
else if (act == 1) ay += 1;
else if (act == 2) ax -= 1;
else ax += 1;
ax = min(max(ax, 0), BOARD - 1);
ay = min(max(ay, 0), BOARD - 1);
int hit = (valid && ax == fx && ay == fy) ? 1 : 0;
if (valid) {
agent[ge * 2 + 0] = ax;
agent[ge * 2 + 1] = ay;
food[ge * 2 + 0] = fx;
food[ge * 2 + 1] = fy;
rng[ge] = r;
rewards[ge] += (float)hit;
#pragma unroll
for (int a = 0; a < NACT; a++) last_logits[ge * NACT + a] = lg[a];
}
// Every block races to set the same flag; a cheap pre-read keeps all but the
// first few from serialising on the L2 atomic (2048 blocks at the largest shape).
unsigned mask = __ballot_sync(0xffffffff, hit);
if (lane == 0 && mask && !__ldcv(&anyhit[step])) atomicOr(&anyhit[step], 1);
}
}
// --------------------------------------------------------------------------------------
// Exact fp32 reference-shaped policy forward (used by the public policy_forward API).
// One block handles PF_E envs; 256 threads = one hidden unit each.
// --------------------------------------------------------------------------------------
#define PF_E 8
__global__ __launch_bounds__(256) void policy_fwd_kernel(
const float *__restrict__ obs, // (N,4)
const float *__restrict__ state, // (N,3,256)
const float *__restrict__ wenc, // (256,4)
const float *__restrict__ benc, // (256,)
const float *__restrict__ wgruT, // (3,256,768) transposed: [layer][k][gaterow]
const float *__restrict__ wa, // (4,256)
const float *__restrict__ ba, // (4,)
const float *__restrict__ wv, // (1,256)
const float *__restrict__ bv, // (1,)
float *__restrict__ logits, // (N,4)
float *__restrict__ nstate, // (N,3,256)
float *__restrict__ value, // (N,)
int N) {
__shared__ float sh[PF_E][HID];
__shared__ float sred[PF_E][NACT + 1][8];
const int j = threadIdx.x;
const int eb = blockIdx.x * PF_E;
float h[PF_E];
#pragma unroll
for (int u = 0; u < PF_E; u++) {
int e = eb + u;
float acc = benc[j];
if (e < N) {
#pragma unroll
for (int i = 0; i < 4; i++) acc += wenc[j * 4 + i] * obs[e * 4 + i];
}
h[u] = acc;
sh[u][j] = acc;
}
__syncthreads();
for (int L = 0; L < 3; L++) {
float zh[PF_E], zg[PF_E], zp[PF_E];
#pragma unroll
for (int u = 0; u < PF_E; u++) { zh[u] = 0.f; zg[u] = 0.f; zp[u] = 0.f; }
const float *wp = wgruT + (size_t)L * HID * NGATE;
for (int k = 0; k < HID; k++) {
float a0 = wp[(size_t)k * NGATE + j];
float a1 = wp[(size_t)k * NGATE + HID + j];
float a2 = wp[(size_t)k * NGATE + 2 * HID + j];
#pragma unroll
for (int u = 0; u < PF_E; u++) {
float hv = sh[u][k];
zh[u] += a0 * hv;
zg[u] += a1 * hv;
zp[u] += a2 * hv;
}
}
__syncthreads();
#pragma unroll
for (int u = 0; u < PF_E; u++) {
int e = eb + u;
float st = (e < N) ? state[((size_t)e * 3 + L) * HID + j] : 0.f;
float o = st + sigmoidf(zg[u]) * (tanhf(zh[u]) - st);
float p = sigmoidf(zp[u]);
h[u] = p * o + (1.0f - p) * h[u];
sh[u][j] = h[u];
if (e < N) nstate[((size_t)e * 3 + L) * HID + j] = o;
}
__syncthreads();
}
// heads
{
float acc[PF_E][NACT + 1];
#pragma unroll
for (int u = 0; u < PF_E; u++)
#pragma unroll
for (int a = 0; a < NACT + 1; a++) acc[u][a] = 0.f;
#pragma unroll
for (int u = 0; u < PF_E; u++) {
#pragma unroll
for (int a = 0; a < NACT; a++) acc[u][a] = wa[a * HID + j] * h[u];
acc[u][NACT] = wv[j] * h[u];
}
const int lane = j & 31, w = j >> 5;
#pragma unroll
for (int u = 0; u < PF_E; u++)
#pragma unroll
for (int a = 0; a < NACT + 1; a++) {
float v = acc[u][a];
#pragma unroll
for (int off = 16; off; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off);
if (lane == 0) sred[u][a][w] = v;
}
__syncthreads();
if (j < PF_E * (NACT + 1)) {
int u = j / (NACT + 1), a = j % (NACT + 1);
float v = 0.f;
#pragma unroll
for (int i = 0; i < 8; i++) v += sred[u][a][i];
int e = eb + u;
if (e < N) {
if (a < NACT) logits[e * NACT + a] = v + ba[a];
else value[e] = v + bv[0];
}
}
}
}
// --------------------------------------------------------------------------------------
// Public env_step (two kernels: move+hit, then respawn)
// --------------------------------------------------------------------------------------
__global__ void env_move(const float *__restrict__ agent, const float *__restrict__ food,
const long long *__restrict__ actions, float *__restrict__ nagent,
float *__restrict__ reward, int *__restrict__ anyflag, int N) {
int e = blockIdx.x * blockDim.x + threadIdx.x;
int hit = 0;
if (e < N) {
float ax = agent[e * 2 + 0], ay = agent[e * 2 + 1];
long long a = actions[e];
if (a == 0) ay -= 1.f;
else if (a == 1) ay += 1.f;
else if (a == 2) ax -= 1.f;
else if (a == 3) ax += 1.f;
ax = fminf(fmaxf(ax, 0.f), (float)(BOARD - 1));
ay = fminf(fmaxf(ay, 0.f), (float)(BOARD - 1));
nagent[e * 2 + 0] = ax; nagent[e * 2 + 1] = ay;
hit = (ax == food[e * 2 + 0] && ay == food[e * 2 + 1]) ? 1 : 0;
reward[e] = (float)hit;
}
unsigned m = __ballot_sync(0xffffffff, hit);
if ((threadIdx.x & 31) == 0 && m) atomicOr(anyflag, 1);
}
__global__ void env_respawn(const float *__restrict__ nagent, float *__restrict__ food,
long long *__restrict__ rng, const int *__restrict__ anyflag, int N) {
int e = blockIdx.x * blockDim.x + threadIdx.x;
if (e >= N) return;
if (!*anyflag) return;
long long r = rng[e];
r = (long long)((unsigned long long)r * LCG_A + 1ULL) & LCG_MASK;
float fx = (float)(r % BOARD);
r = (long long)((unsigned long long)r * LCG_A + 1ULL) & LCG_MASK;
float fy = (float)(r % BOARD);
rng[e] = r;
if (nagent[e * 2 + 0] == food[e * 2 + 0] && nagent[e * 2 + 1] == food[e * 2 + 1]) {
food[e * 2 + 0] = fx; food[e * 2 + 1] = fy;
}
}
// --------------------------------------------------------------------------------------
// host entry points
// --------------------------------------------------------------------------------------
torch::Tensor prep_wfrag(torch::Tensor w) { // w: (L,768,256) contiguous fp32 cuda
int nl = (int)w.size(0);
auto out = torch::empty({nl * 16 * 96 * 32 * 4}, w.options().dtype(torch::kInt32));
int total = nl * 16 * 96 * 32;
permute_weights<<<(total + 255) / 256, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
w.data_ptr<float>(), reinterpret_cast<uint4 *>(out.data_ptr<int>()), total);
return out;
}
torch::Tensor prep_l0(torch::Tensor wg0, torch::Tensor wenc, torch::Tensor benc) {
auto stream = at::cuda::getCurrentCUDAStream();
auto M0 = torch::empty({NGATE, 4}, wg0.options());
auto c0 = torch::empty({NGATE}, wg0.options());
fold_l0<<<NGATE, 256, 0, stream>>>(wg0.data_ptr<float>(), wenc.data_ptr<float>(),
benc.data_ptr<float>(), M0.data_ptr<float>(),
c0.data_ptr<float>());
auto out = torch::empty({128 * 32 * 4}, wg0.options().dtype(torch::kInt32));
permute_l0<<<(128 * 32 + 255) / 256, 256, 0, stream>>>(
M0.data_ptr<float>(), c0.data_ptr<float>(), wenc.data_ptr<float>(),
benc.data_ptr<float>(), reinterpret_cast<uint4 *>(out.data_ptr<int>()));
return out;
}
void rollout(torch::Tensor wfrag, torch::Tensor wfrag0, torch::Tensor wa, torch::Tensor ba,
torch::Tensor state, torch::Tensor agent, torch::Tensor food, torch::Tensor rng,
torch::Tensor rewards, torch::Tensor last_logits, torch::Tensor anyhit, int64_t N,
int64_t horizon) {
int nblk = (int)((N + RENV - 1) / RENV);
size_t sb = smem_bytes();
static bool attr_set = false;
if (!attr_set) {
cudaFuncSetAttribute(rollout_step, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)sb);
cudaFuncSetAttribute(rollout_step, cudaFuncAttributePreferredSharedMemoryCarveout, 100);
attr_set = true;
}
auto stream = at::cuda::getCurrentCUDAStream();
for (int t = 0; t < horizon; t++) {
rollout_step<<<nblk, NTHREADS, sb, stream>>>(
reinterpret_cast<const uint4 *>(wfrag.data_ptr<int>()),
reinterpret_cast<const uint4 *>(wfrag0.data_ptr<int>()), wa.data_ptr<float>(),
ba.data_ptr<float>(), state.data_ptr<float>(), agent.data_ptr<int>(),
food.data_ptr<int>(), reinterpret_cast<long long *>(rng.data_ptr<int64_t>()),
rewards.data_ptr<float>(), last_logits.data_ptr<float>(), anyhit.data_ptr<int>(),
(int)N, t);
}
}
std::vector<torch::Tensor> policy_forward_cuda(torch::Tensor obs, torch::Tensor state,
torch::Tensor wenc, torch::Tensor benc,
torch::Tensor wgruT, torch::Tensor wa,
torch::Tensor ba, torch::Tensor wv,
torch::Tensor bv) {
int N = obs.size(0);
auto logits = torch::empty({N, NACT}, obs.options());
auto nstate = torch::empty({N, 3, HID}, obs.options());
auto value = torch::empty({N}, obs.options());
int nblk = (N + PF_E - 1) / PF_E;
policy_fwd_kernel<<<nblk, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
obs.data_ptr<float>(), state.data_ptr<float>(), wenc.data_ptr<float>(),
benc.data_ptr<float>(), wgruT.data_ptr<float>(), wa.data_ptr<float>(), ba.data_ptr<float>(),
wv.data_ptr<float>(), bv.data_ptr<float>(), logits.data_ptr<float>(),
nstate.data_ptr<float>(), value.data_ptr<float>(), N);
return {logits, nstate, value};
}
std::vector<torch::Tensor> env_step_cuda(torch::Tensor agent, torch::Tensor food,
torch::Tensor actions, torch::Tensor rng) {
int N = agent.size(0);
auto nagent = torch::empty_like(agent);
auto nfood = food.clone();
auto nrng = rng.clone();
auto reward = torch::empty({N}, agent.options());
auto flag = torch::zeros({1}, agent.options().dtype(torch::kInt32));
auto stream = at::cuda::getCurrentCUDAStream();
int nb = (N + 255) / 256;
env_move<<<nb, 256, 0, stream>>>(agent.data_ptr<float>(), food.data_ptr<float>(),
reinterpret_cast<const long long *>(actions.data_ptr<int64_t>()), nagent.data_ptr<float>(),
reward.data_ptr<float>(), flag.data_ptr<int>(), N);
env_respawn<<<nb, 256, 0, stream>>>(nagent.data_ptr<float>(), nfood.data_ptr<float>(),
reinterpret_cast<long long *>(nrng.data_ptr<int64_t>()), flag.data_ptr<int>(), N);
return {nagent, nfood, reward, nrng};
}
// --------------------------------------------------------------------------------------
// Host-side mt19937 matching at::CPUGeneratorImpl (init_genrand seeding, 32-bit draws).
// torch.randint(0, 11) is `engine() % 11`, so reproducing the raw stream reproduces the
// tensor bit-for-bit; doing it here (vectorised block twist) is ~4x faster than the
// generic per-element ATen path and this sits directly in run()'s critical path.
// --------------------------------------------------------------------------------------
namespace mtgen {
constexpr int N_ = 624, M_ = 397;
constexpr uint32_t MATRIX_A = 0x9908b0dfu, UMASK = 0x80000000u, LMASK = 0x7fffffffu;
static inline uint32_t twist1(uint32_t u, uint32_t v) {
return (((u & UMASK) | (v & LMASK)) >> 1) ^ ((v & 1u) ? MATRIX_A : 0u);
}
static inline uint32_t temper1(uint32_t y) {
y ^= (y >> 11);
y ^= (y << 7) & 0x9d2c5680u;
y ^= (y << 15) & 0xefc60000u;
y ^= (y >> 18);
return y;
}
static void next_state_scalar(uint32_t *st) {
uint32_t *p = st;
for (int j = N_ - M_ + 1; --j; p++) *p = p[M_] ^ twist1(p[0], p[1]);
for (int j = M_; --j; p++) *p = p[M_ - N_] ^ twist1(p[0], p[1]);
*p = p[M_ - N_] ^ twist1(p[0], st[0]);
}
#if defined(__x86_64__)
__attribute__((target("avx2"))) static inline __m256i twist8(__m256i u, __m256i v) {
const __m256i um = _mm256_set1_epi32((int)UMASK), lm = _mm256_set1_epi32((int)LMASK);
__m256i mixed = _mm256_or_si256(_mm256_and_si256(u, um), _mm256_and_si256(v, lm));
__m256i odd = _mm256_and_si256(v, _mm256_set1_epi32(1));
__m256i mask = _mm256_sub_epi32(_mm256_setzero_si256(), odd);
return _mm256_xor_si256(_mm256_srli_epi32(mixed, 1),
_mm256_and_si256(mask, _mm256_set1_epi32((int)MATRIX_A)));
}
__attribute__((target("avx2"))) static void next_state_avx2(uint32_t *st) {
int i = 0;
for (; i + 8 <= N_ - M_; i += 8) { // reads are all pre-twist here
__m256i u = _mm256_loadu_si256((const __m256i *)(st + i));
__m256i v = _mm256_loadu_si256((const __m256i *)(st + i + 1));
__m256i w = _mm256_loadu_si256((const __m256i *)(st + i + M_));
_mm256_storeu_si256((__m256i *)(st + i), _mm256_xor_si256(w, twist8(u, v)));
}
for (; i < N_ - M_; i++) st[i] = st[i + M_] ^ twist1(st[i], st[i + 1]);
for (; i + 8 <= N_ - 1; i += 8) { // st[i-227] is already updated, distance 227 >= 8
__m256i u = _mm256_loadu_si256((const __m256i *)(st + i));
__m256i v = _mm256_loadu_si256((const __m256i *)(st + i + 1));
__m256i w = _mm256_loadu_si256((const __m256i *)(st + i + M_ - N_));
_mm256_storeu_si256((__m256i *)(st + i), _mm256_xor_si256(w, twist8(u, v)));
}
for (; i < N_ - 1; i++) st[i] = st[i + M_ - N_] ^ twist1(st[i], st[i + 1]);
st[N_ - 1] = st[M_ - 1] ^ twist1(st[N_ - 1], st[0]);
}
__attribute__((target("avx2"))) static void temper_mod_avx2(const uint32_t *st, int32_t *out,
int cnt, uint32_t range) {
int i = 0;
const __m256i c7 = _mm256_set1_epi32((int)0x9d2c5680u), c15 = _mm256_set1_epi32((int)0xefc60000u);
for (; i + 8 <= cnt; i += 8) {
__m256i y = _mm256_loadu_si256((const __m256i *)(st + i));
y = _mm256_xor_si256(y, _mm256_srli_epi32(y, 11));
y = _mm256_xor_si256(y, _mm256_and_si256(_mm256_slli_epi32(y, 7), c7));
y = _mm256_xor_si256(y, _mm256_and_si256(_mm256_slli_epi32(y, 15), c15));
y = _mm256_xor_si256(y, _mm256_srli_epi32(y, 18));
// y % 11 via the exact magic multiply q = (y * 3123612579) >> 35
const __m256i mg = _mm256_set1_epi64x(3123612579LL);
__m256i qe = _mm256_srli_epi64(_mm256_mul_epu32(y, mg), 35);
__m256i qo = _mm256_srli_epi64(_mm256_mul_epu32(_mm256_srli_epi64(y, 32), mg), 35);
__m256i q = _mm256_blend_epi32(qe, _mm256_slli_epi64(qo, 32), 0xAA);
__m256i r = _mm256_sub_epi32(y, _mm256_mullo_epi32(q, _mm256_set1_epi32((int)range)));
_mm256_storeu_si256((__m256i *)(out + i), r);
}
for (; i < cnt; i++) out[i] = (int32_t)(temper1(st[i]) % range);
}
static bool have_avx2() {
static int c = -1;
if (c < 0) c = __builtin_cpu_supports("avx2") ? 1 : 0;
return c == 1;
}
#endif
static void randint_fill(int32_t *out, int64_t n, uint64_t seed, uint32_t range) {
uint32_t st[N_];
st[0] = (uint32_t)(seed & 0xffffffffu);
for (int j = 1; j < N_; j++)
st[j] = 1812433253u * (st[j - 1] ^ (st[j - 1] >> 30)) + (uint32_t)j;
int64_t done = 0;
while (done < n) {
#if defined(__x86_64__)
if (have_avx2()) next_state_avx2(st); else next_state_scalar(st);
#else
next_state_scalar(st);
#endif
int cnt = (int)std::min<int64_t>(N_, n - done);
#if defined(__x86_64__)
if (have_avx2() && range == 11) temper_mod_avx2(st, out + done, cnt, range);
else
#endif
for (int i = 0; i < cnt; i++) out[done + i] = (int32_t)(temper1(st[i]) % range);
done += cnt;
}
}
} // namespace mtgen
torch::Tensor randint_cpu(int64_t n, int64_t seed, int64_t range) {
auto opts = torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU).pinned_memory(true);
auto t = torch::empty({n}, opts);
mtgen::randint_fill(t.data_ptr<int32_t>(), n, (uint64_t)seed, (uint32_t)range);
return t;
}
// One fused setup kernel: splits the uploaded rng draws into agent/food, seeds the
// per-env LCG state and clears the accumulators, so run() issues one launch here
// instead of half a dozen ATen ops.
__global__ void init_run_kernel(const int *__restrict__ draws, int *__restrict__ agent,
int *__restrict__ food, long long *__restrict__ rng,
float *__restrict__ rewards, float *__restrict__ last_logits,
int N, long long rng_base) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N) return;
agent[i * 2 + 0] = draws[i * 2 + 0];
agent[i * 2 + 1] = draws[i * 2 + 1];
food[i * 2 + 0] = draws[(N + i) * 2 + 0];
food[i * 2 + 1] = draws[(N + i) * 2 + 1];
rng[i] = rng_base + i;
rewards[i] = 0.0f;
last_logits[i * 4 + 0] = 0.0f;
last_logits[i * 4 + 1] = 0.0f;
last_logits[i * 4 + 2] = 0.0f;
last_logits[i * 4 + 3] = 0.0f;
}
std::vector<torch::Tensor> init_run(torch::Tensor draws, int64_t N, int64_t rng_base) {
auto iopt = draws.options();
auto fopt = draws.options().dtype(torch::kFloat32);
auto agent = torch::empty({N, 2}, iopt);
auto food = torch::empty({N, 2}, iopt);
auto rng = torch::empty({N}, draws.options().dtype(torch::kInt64));
auto rewards = torch::empty({N}, fopt);
auto last_logits = torch::empty({N, NACT}, fopt);
init_run_kernel<<<(int)((N + 255) / 256), 256, 0, at::cuda::getCurrentCUDAStream()>>>(
draws.data_ptr<int>(), agent.data_ptr<int>(), food.data_ptr<int>(),
reinterpret_cast<long long *>(rng.data_ptr<int64_t>()), rewards.data_ptr<float>(),
last_logits.data_ptr<float>(), (int)N, (long long)rng_base);
return {agent, food, rng, rewards, last_logits};
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("prep_wfrag", &prep_wfrag);
m.def("randint_cpu", &randint_cpu);
m.def("init_run", &init_run);
m.def("prep_l0", &prep_l0);
m.def("rollout", &rollout);
m.def("policy_forward_cuda", &policy_forward_cuda);
m.def("env_step_cuda", &env_step_cuda);
}
"""
def _build():
import hashlib
import shutil
import sys
from torch.utils.cpp_extension import _get_build_directory, load
if shutil.which("ninja") is None: # venv installs land next to the interpreter
os.environ["PATH"] = str(Path(sys.executable).parent) + os.pathsep + os.environ["PATH"]
cc = torch.cuda.get_device_capability(0)
arch = f"{cc[0]}{cc[1]}"
tag = hashlib.sha1(_CUDA_SRC.encode()).hexdigest()[:10]
name = f"grid_mingru_sps_{arch}_{tag}"
build_dir = Path(_get_build_directory(name, verbose=False))
build_dir.mkdir(parents=True, exist_ok=True)
src = build_dir / "kernels.cu"
if not src.exists() or src.read_text() != _CUDA_SRC:
src.write_text(_CUDA_SRC)
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", f"{cc[0]}.{cc[1]}")
def _try(sfx):
return load(
name=name,
sources=[str(src)],
extra_cuda_cflags=[
"-O3",
f"-gencode=arch=compute_{arch}{sfx},code=sm_{arch}{sfx}",
"-DNDEBUG",
],
extra_cflags=["-O3"],
verbose=False,
)
try: # sm_90a/100a/120a expose the arch-specific instruction set
return _try("a" if cc[0] >= 9 else "")
except Exception:
return _try("")
_EXT = None
def _ext():
global _EXT
if _EXT is None:
_EXT = _build()
return _EXT
class Model(nn.Module):
def __init__(self):
super().__init__()
self.w_enc = nn.Parameter(torch.empty(HIDDEN, OBS_DIM))
self.b_enc = nn.Parameter(torch.zeros(HIDDEN))
self.w_gru = nn.Parameter(torch.empty(GRU_LAYERS, GRU_OUT, HIDDEN))
self.w_a = nn.Parameter(torch.empty(NUM_ACTIONS, HIDDEN))
self.b_a = nn.Parameter(torch.zeros(NUM_ACTIONS))
self.w_v = nn.Parameter(torch.empty(1, HIDDEN))
self.b_v = nn.Parameter(torch.zeros(1))
self.reset_parameters(0)
def reset_parameters(self, seed: int = 0) -> None:
g = torch.Generator(device="cpu")
g.manual_seed(seed)
for p in self.parameters():
tmp = torch.empty(p.shape, dtype=p.dtype, device="cpu")
tmp.normal_(0.0, 0.02, generator=g)
p.data.copy_(tmp)
def forward(self, obs: torch.Tensor, state: torch.Tensor):
return policy_forward(self, obs, state)
# ----------------------------------------------------------------------------------
# public single-step APIs (exact fp32 paths)
# ----------------------------------------------------------------------------------
def policy_forward(model, obs: torch.Tensor, state: torch.Tensor):
ext = _ext()
wgruT = model.w_gru.detach().transpose(1, 2).contiguous()
logits, nstate, value = ext.policy_forward_cuda(
obs.contiguous().float(),
state.contiguous().float(),
model.w_enc.detach().contiguous().float(),
model.b_enc.detach().contiguous().float(),
wgruT.float(),
model.w_a.detach().contiguous().float(),
model.b_a.detach().contiguous().float(),
model.w_v.detach().contiguous().float(),
model.b_v.detach().contiguous().float(),
)
return logits, nstate, value
def env_step(agent: torch.Tensor, food: torch.Tensor, actions: torch.Tensor,
rng_state: torch.Tensor):
ext = _ext()
a, f, r, g = ext.env_step_cuda(
agent.contiguous().float(),
food.contiguous().float(),
actions.contiguous().to(torch.int64),
rng_state.contiguous().to(torch.int64),
)
return a.to(agent.dtype), f.to(food.dtype), r, g
# ----------------------------------------------------------------------------------
# rollout
# ----------------------------------------------------------------------------------
def _prepare(model):
ext = _ext()
w_gru = model.w_gru.detach()
wfrag0 = ext.prep_l0(w_gru[0].contiguous(), model.w_enc.detach().contiguous(),
model.b_enc.detach().contiguous())
wfrag = ext.prep_wfrag(w_gru[1:].contiguous())
return (wfrag, wfrag0, model.w_a.detach().contiguous(), model.b_a.detach().contiguous())
def run(num_envs: int, horizon: int, seed: int, model=None) -> dict:
ext = _ext()
device = torch.device("cuda:0")
if model is None:
model = Model()
if model.w_gru.device != device:
model = model.to(device)
if num_envs <= 0:
z = torch.zeros(0, device=device)
return {"rewards": z, "positions": torch.zeros(0, 2, dtype=torch.int64, device=device),
"last_logits": torch.zeros(0, NUM_ACTIONS, device=device),
"state": torch.zeros(0, GRU_LAYERS, HIDDEN, device=device)}
wfrag, wfrag0, wa, ba = _prepare(model)
# Reference draws agent then food from a fresh cpu mt19937; randint_cpu reproduces
# that 32-bit stream exactly (int32/int64 consume identical draws for range < 2^32)
# into pinned memory, ~6x faster than the generic ATen path.
draws = ext.randint_cpu(4 * num_envs, seed, BOARD).to(device, non_blocking=True)
agent, food, rng_state, rewards, last_logits = ext.init_run(draws, num_envs, seed * 10007)
state = torch.zeros(num_envs, GRU_LAYERS, HIDDEN, device=device)
anyhit = torch.zeros(max(horizon, 1), device=device, dtype=torch.int32)
ext.rollout(wfrag, wfrag0, wa, ba, state, agent, food, rng_state,
rewards, last_logits, anyhit, num_envs, horizon)
return {
"rewards": rewards,
"positions": agent.to(torch.int64),
"last_logits": last_logits,
"state": state,
}
def get_init_inputs():
return []
def get_inputs():
return []
_ext() # compile at import time, not inside a timed run()
20260721_144141_claude_claude-opus-4-8_04_grid_mingru_sps