"""Grid-foraging PPO training megakernel. The entire PPO training run -- rollout (all HORIZON steps x all envs, policy forward, action sampling, env step, reward, GAE), advantage normalization, and the full PPO update (4 epochs x 4 minibatches of backprop + grad-norm clip + Adam) -- executes inside ONE persistent cooperative CUDA kernel launch per training run. Grid-wide phase synchronization uses cooperative-groups grid.sync(); kernel-launch count is O(1) per run and never scales with env steps, horizon, epochs, or minibatches. No CUDA graphs, no torch.compile, no prebuilt RL libraries: the environment, policy, and PPO update are hand-written fused CUDA. There is no Python-side per-step or per-minibatch kernel loop anywhere: train() issues exactly one cooperative launch and one device->host copy of the 40-iteration return curve. Parallel decomposition (the kernel body is a template instantiated at both NT=512 and NT=256; the host picks the fattest variant that actually fits with a cooperative co-resident grid on the current arch): * NT threads (NT/32 warps) per block, one block per SM resident (grid capped by co-residency so the cooperative launch is always legal). * Rollout: each warp owns whole environments (balanced contiguous ranges). The 32 steps of an env run lane-parallel inside the warp (lane l computes hidden units l and l+32; a butterfly shuffle reduction yields the 4 actor logits + value), while every lane redundantly - hence bit-identically - tracks the env state and PCG RNG. GAE is a per-warp backward recurrence over a shared-memory reward/value stash. * Advantage mean/std: per-warp partials -> block -> global double atomics; normalization is applied on the fly when the update consumes samples. * PPO update: 16 minibatch barriers per iteration. For perfect balance, each iteration first scatters rollout slots into 16 per-(epoch,minibatch) index lists (membership = hash of the slot index; uniform partition, each sample used exactly once per epoch), built between rollout and the update while rollout stragglers finish. Warps then take contiguous, count-balanced list slices and process samples warp-collectively two-at-a-time with software pipelined gathers. Analytic PPO gradients (exactly matching autograd of loss = pg + 0.5*vloss - 0.01*entropy) accumulate in per-lane registers (each lane owns its own hidden units' gradient slots), then per-block shared-memory sums, then pre-zeroed global atomic slots per minibatch. After grid.sync every block performs the identical grad-norm clip + Adam update on its own shared-memory copy of the 645 parameters - redundantly deterministic, so all replicas stay bitwise identical with no weight broadcast. """ from __future__ import annotations import os import torch OP_TYPE = "rl_grid_ppo" HARDWARE_REQUIRED = ["RTX_PRO_6000"] # --- Task constants (must match reference.py exactly) ------------------------ GRID = 11 NUM_ENVS = 4096 HORIZON = 32 ROLLOUT = 32 OBS_DIM = 4 N_ACT = 4 HIDDEN = 64 GAMMA = 0.99 LAM = 0.95 CLIP = 0.2 EPOCHS = 4 MINIBATCHES = 4 LR = 3.0e-3 ENT_COEF = 0.01 VF_COEF = 0.5 MAX_GRAD_NORM = 0.5 N_SAMPLES = ROLLOUT * NUM_ENVS N_MINIBATCHES = EPOCHS * MINIBATCHES NPARAM = 645 # 4*64 + 64 + 4*64 + 4 + 64 + 1 _CUDA_SRC = r""" #include #include #include #include #include #include namespace cg = cooperative_groups; #define NUM_ENVS 4096 #define T_STEPS 32 #define N_SAMPLES 131072 #define NPARAM 645 #define N_MB 16 #define INV_MB (1.0f / 32768.0f) #define GAMMA 0.99f #define GL (0.99f * 0.95f) #define LRF 3.0e-3f #define ENT_COEF 0.01f #define MAXN 0.5f #define A_EPS 1e-8f #define BETA1 0.9f #define BETA2 0.999f #define INV_N (1.0 / 131072.0) // Parameter layout (645 floats, shared-memory master copy per block): // [0,256) W1T : k*64 + j (obs k -> hidden j), init U(-1/2,1/2) // [256,320) B1 : j zero // [320,576) W2 : m*64 + j (hidden j -> logit m), init U(-1/8,1/8) // [576,580) B2 : m zero // [580,644) W3 : j (hidden j -> value), init U(-1/8,1/8) // [644] B3 zero template struct SharedMem { float w[NPARAM]; float am[NPARAM]; float av[NPARAM]; float gacc[NPARAM]; float rew_stash[NWARP][T_STEPS]; float val_stash[NWARP][T_STEPS]; double stats_blk[2]; float curve_blk; float red[NWARP]; }; __device__ __forceinline__ uint64_t splitmix64(uint64_t x) { x += 0x9E3779B97F4A7C15ULL; x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ULL; x = (x ^ (x >> 27)) * 0x94D049BB133111EBULL; return x ^ (x >> 31); } __device__ __forceinline__ float u01(uint64_t h) { return (float)(h >> 40) * (1.0f / 16777216.0f); } // PCG32 (XSH-RR) per-env RNG; advanced redundantly & identically by all lanes. struct Rng { uint64_t s; }; __device__ __forceinline__ uint32_t rng_next(Rng* r) { uint64_t old = r->s; r->s = old * 6364136223846793005ULL + 1442695040888963407ULL; uint32_t x = (uint32_t)(((old >> 18) ^ old) >> 27); uint32_t rot = (uint32_t)(old >> 59); return (x >> rot) | (x << ((32u - rot) & 31u)); } __device__ __forceinline__ float rng_u01(Rng* r) { return (float)(rng_next(r) >> 8) * (1.0f / 16777216.0f); } __device__ __forceinline__ int rng_cell(Rng* r) { return (int)(((uint64_t)rng_next(r) * 11ULL) >> 32); // uniform in [0,10] } // Lane-parallel policy forward. Lane l owns hidden units l and l+32. // Fills out5[0..3]=logits, out5[4]=value in every lane; returns h1,h2. __device__ __forceinline__ void policy_fwd(const float* w, int lane, const float obs[4], float* h1o, float* h2o, float* out5) { float a1 = w[256 + lane], a2 = w[256 + 32 + lane]; #pragma unroll for (int k = 0; k < 4; ++k) { a1 += w[k * 64 + lane] * obs[k]; a2 += w[k * 64 + 32 + lane] * obs[k]; } float h1 = __tanhf(a1); float h2 = __tanhf(a2); *h1o = h1; *h2o = h2; #pragma unroll for (int m = 0; m < 4; ++m) { out5[m] = w[320 + m * 64 + lane] * h1 + w[320 + m * 64 + 32 + lane] * h2; } out5[4] = w[580 + lane] * h1 + w[580 + 32 + lane] * h2; #pragma unroll for (int off = 16; off >= 1; off >>= 1) { #pragma unroll for (int m = 0; m < 5; ++m) { out5[m] += __shfl_xor_sync(0xffffffffu, out5[m], off); } } #pragma unroll for (int m = 0; m < 4; ++m) out5[m] += w[576 + m]; out5[4] += w[644]; } // grad wrt hidden unit j given head weights and upstream head grads (dz, dv) __device__ __forceinline__ float w_dot_head(const float* w, int j, const float dz[4], float dv) { return w[320 + j] * dz[0] + w[320 + 64 + j] * dz[1] + w[320 + 128 + j] * dz[2] + w[320 + 192 + j] * dz[3] + w[580 + j] * dv; } template __device__ __forceinline__ void ppo_body(float4* __restrict__ obsQ, // [N_SAMPLES] observations float4* __restrict__ metaQ, // [N_SAMPLES] (logp_old, act, adv, ret) float* __restrict__ curveG, // [iters] reward sums (unscaled) double* __restrict__ statsG, // [2] sum, sumsq of advantages float* __restrict__ gslotG, // [iters*16*645] grad slots int* __restrict__ listG, // [16*40960] minibatch member lists int* __restrict__ cntG, // [iters*16] per-minibatch counts int iters, uint64_t seed) { cg::grid_group grid = cg::this_grid(); __shared__ SharedMem smobj; SharedMem* sm = &smobj; __shared__ float gvals[NPARAM]; const int tid = threadIdx.x; const int lane = tid & 31; const int wid = tid >> 5; const int nb = gridDim.x; const int gw = blockIdx.x * NWARP + wid; // global warp id const int TW = nb * NWARP; // total warps // ---- weight + Adam-state init (identical in every block) ---------------- for (int k = tid; k < NPARAM; k += NT) { float val; if (k < 256) { val = u01(splitmix64(seed ^ (0x1A2B3C4D00000000ULL + (uint64_t)k))) - 0.5f; } else if (k < 320) { val = 0.0f; } else if (k < 576) { val = (u01(splitmix64(seed ^ (0x5E6F7A8B00000000ULL + (uint64_t)k))) - 0.5f) * 0.25f; } else if (k < 580) { val = 0.0f; } else if (k < 644) { val = (u01(splitmix64(seed ^ (0x9C0D1E2F00000000ULL + (uint64_t)k))) - 0.5f) * 0.25f; } else { val = 0.0f; } sm->w[k] = val; sm->am[k] = 0.0f; sm->av[k] = 0.0f; } __syncthreads(); for (int it = 0; it < iters; ++it) { // ========================= ROLLOUT + GAE =========================== if (tid == 0) { sm->stats_blk[0] = 0.0; sm->stats_blk[1] = 0.0; sm->curve_blk = 0.0f; } __syncthreads(); const int e_lo = (int)(((long long)gw * NUM_ENVS) / TW); const int e_hi = (int)(((long long)(gw + 1) * NUM_ENVS) / TW); for (int e = e_lo; e < e_hi; ++e) { Rng rng; rng.s = splitmix64((seed * 0x6A09E667F3BCC909ULL) ^ ((uint64_t)(e + 1) * 0x9E3779B97F4A7C15ULL) ^ ((uint64_t)(it + 1) * 0x165667B19E3779F9ULL)); int ax = rng_cell(&rng), ay = rng_cell(&rng); int fx = rng_cell(&rng), fy = rng_cell(&rng); float env_ret = 0.0f; #pragma unroll 1 for (int t = 0; t < T_STEPS; ++t) { float obs[4] = { (float)(fx - ax) * (1.0f / 11.0f), (float)(fy - ay) * (1.0f / 11.0f), (float)ax * 0.1f, (float)ay * 0.1f }; float h1, h2, out5[5]; policy_fwd(sm->w, lane, obs, &h1, &h2, out5); // categorical sample from softmax(logits) float mx = fmaxf(fmaxf(out5[0], out5[1]), fmaxf(out5[2], out5[3])); float ex0 = __expf(out5[0] - mx), ex1 = __expf(out5[1] - mx); float ex2 = __expf(out5[2] - mx), ex3 = __expf(out5[3] - mx); float Z = ex0 + ex1 + ex2 + ex3; float invZ = __frcp_rn(Z); float q0 = ex0 * invZ, q1 = ex1 * invZ, q2 = ex2 * invZ; float u = rng_u01(&rng); int act = (u < q0) ? 0 : (u < q0 + q1) ? 1 : (u < q0 + q1 + q2) ? 2 : 3; float lz = __logf(Z); float sel = (act == 0) ? out5[0] : (act == 1) ? out5[1] : (act == 2) ? out5[2] : out5[3]; float logp = (sel - mx) - lz; const int s = t * NUM_ENVS + e; if (lane == 0) obsQ[s] = make_float4(obs[0], obs[1], obs[2], obs[3]); else if (lane == 1) metaQ[s].x = logp; else if (lane == 2) metaQ[s].y = (float)act; // env step: 0=up,1=down,2=left,3=right, clamped to [0,10] if (act == 0) ay = (ay > 0) ? ay - 1 : 0; else if (act == 1) ay = (ay < 10) ? ay + 1 : 10; else if (act == 2) ax = (ax > 0) ? ax - 1 : 0; else ax = (ax < 10) ? ax + 1 : 10; float r = (ax == fx && ay == fy) ? 1.0f : 0.0f; if (r > 0.0f) { fx = rng_cell(&rng); fy = rng_cell(&rng); } env_ret += r; __syncwarp(); if (lane == 3) sm->rew_stash[wid][t] = r; if (lane == 4) sm->val_stash[wid][t] = out5[4]; } __syncwarp(); // ---- GAE backward recurrence (all lanes redundantly) ----------- float gae = 0.0f; float adv_sum = 0.0f, adv_sumsq = 0.0f; #pragma unroll 1 for (int t = T_STEPS - 1; t >= 0; --t) { float v_t = sm->val_stash[wid][t]; float r_t = sm->rew_stash[wid][t]; float delta; if (t == T_STEPS - 1) { delta = r_t - v_t; // no bootstrap past horizon gae = delta; } else { delta = r_t + GAMMA * sm->val_stash[wid][t + 1] - v_t; gae = delta + GL * gae; } const int s = t * NUM_ENVS + e; if ((t & 31) == lane) { metaQ[s].z = gae; metaQ[s].w = gae + v_t; } adv_sum += gae; adv_sumsq += gae * gae; } if (lane == 0) { atomicAdd(&sm->stats_blk[0], (double)adv_sum); atomicAdd(&sm->stats_blk[1], (double)adv_sumsq); atomicAdd(&sm->curve_blk, env_ret); } __syncwarp(); } // ---- build per-minibatch member index lists (balance for update) ---- { __shared__ int cnt16[16], cur16[16]; for (int i = tid; i < 16; i += NT) cnt16[i] = 0; __syncthreads(); const int spb = (N_SAMPLES + nb - 1) / nb; const int s_lo = blockIdx.x * spb; const int s_hi = min(N_SAMPLES, s_lo + spb); const uint64_t lkey = splitmix64(seed ^ (0xC6A4A7935BD1E995ULL + (uint64_t)it * 0x9E3779B97F4A7C15ULL)); for (int s0 = s_lo + (wid << 5); s0 < s_hi; s0 += (NWARP << 5)) { const int ss = s0 + lane; const uint64_t h = (ss < s_hi) ? splitmix64((uint64_t)ss ^ lkey) : 0ULL; const bool ok = (ss < s_hi); #pragma unroll for (int ep = 0; ep < 4; ++ep) { const int mbv = (int)((h >> (56 + 2 * ep)) & 3ULL); #pragma unroll for (int m = 0; m < 4; ++m) { const unsigned b = __ballot_sync(0xffffffffu, ok && mbv == m); if (lane == 0 && b) atomicAdd(&cnt16[ep * 4 + m], __popc(b)); } } } __syncthreads(); if (tid < 16) { cur16[tid] = atomicAdd(&cntG[it * 16 + tid], cnt16[tid]); } __syncthreads(); for (int s0 = s_lo + (wid << 5); s0 < s_hi; s0 += (NWARP << 5)) { const int ss = s0 + lane; const uint64_t h = (ss < s_hi) ? splitmix64((uint64_t)ss ^ lkey) : 0ULL; const bool ok = (ss < s_hi); #pragma unroll for (int ep = 0; ep < 4; ++ep) { const int mbv = (int)((h >> (56 + 2 * ep)) & 3ULL); #pragma unroll for (int m = 0; m < 4; ++m) { const unsigned b = __ballot_sync(0xffffffffu, ok && mbv == m); const int pc = __popc(b); if (pc) { int pos0 = 0; if (lane == 0) pos0 = atomicAdd(&cur16[ep * 4 + m], pc); pos0 = __shfl_sync(0xffffffffu, pos0, 0); if (ok && mbv == m) { const int r = __popc(b & ((1u << lane) - 1u)); listG[(ep * 4 + m) * 40960 + pos0 + r] = ss; } } } } } __syncthreads(); } __syncthreads(); if (tid == 0) { double* st = statsG + 2 * it; atomicAdd(&st[0], sm->stats_blk[0]); atomicAdd(&st[1], sm->stats_blk[1]); atomicAdd(&curveG[it], sm->curve_blk); } grid.sync(); const double mean = statsG[2 * it] * INV_N; double var = statsG[2 * it + 1] * INV_N - mean * mean; var = var * ((double)N_SAMPLES / (double)(N_SAMPLES - 1)); // unbiased std const float inv_std = (float)(1.0 / (sqrt(var) + 1e-8)); const float fmean = (float)mean; // ========================= PPO UPDATE ============================== for (int epc = 0; epc < 4; ++epc) { for (int mb = 0; mb < 4; ++mb) { const int midx = epc * 4 + mb; const int slot0 = (it * N_MB + midx) * NPARAM; for (int k = tid; k < NPARAM; k += NT) sm->gacc[k] = 0.0f; __syncthreads(); // per-lane gradient accumulators (lane owns its units' slots) float a_w1[8], a_w2[8], a_b1[2] = {0.0f, 0.0f}; float a_b2 = 0.0f, a_w3[2] = {0.0f, 0.0f}, a_b3 = 0.0f; #pragma unroll for (int i = 0; i < 8; ++i) { a_w1[i] = 0.0f; a_w2[i] = 0.0f; } // balanced list-driven processing; 2 samples in flight const int mcnt = cntG[it * 16 + midx]; const int* mlist = listG + midx * 40960; const int ilo = (int)(((long long)mcnt * gw) / TW); const int ihi = (int)(((long long)mcnt * (gw + 1)) / TW); int ps0 = (ilo < ihi) ? mlist[ilo] : -1; int ps1 = (ilo + 1 < ihi) ? mlist[ilo + 1] : -1; #pragma unroll 1 for (int i0 = ilo; i0 < ihi; i0 += 2) { const int ns0 = (i0 + 2 < ihi) ? mlist[i0 + 2] : -1; const int ns1 = (i0 + 3 < ihi) ? mlist[i0 + 3] : -1; float4 ob[2], mt[2]; float valid[2] = {0.0f, 0.0f}; #pragma unroll for (int u = 0; u < 2; ++u) { const int s = u ? ps1 : ps0; if (s >= 0) { ob[u] = obsQ[s]; mt[u] = metaQ[s]; valid[u] = 1.0f; } else { ob[u] = make_float4(0.f, 0.f, 0.f, 0.f); mt[u] = make_float4(0.f, 0.f, 0.f, 0.f); } } ps0 = ns0; ps1 = ns1; float obs[2][4], out5[2][5], h1[2], h2[2]; #pragma unroll for (int u = 0; u < 2; ++u) { obs[u][0] = ob[u].x; obs[u][1] = ob[u].y; obs[u][2] = ob[u].z; obs[u][3] = ob[u].w; policy_fwd(sm->w, lane, obs[u], &h1[u], &h2[u], out5[u]); } #pragma unroll for (int u = 0; u < 2; ++u) { const int act = (int)mt[u].y; const float advn = ((mt[u].z - fmean) * inv_std) * valid[u]; float mx = fmaxf(fmaxf(out5[u][0], out5[u][1]), fmaxf(out5[u][2], out5[u][3])); float ex0 = __expf(out5[u][0] - mx), ex1 = __expf(out5[u][1] - mx); float ex2 = __expf(out5[u][2] - mx), ex3 = __expf(out5[u][3] - mx); float Z = ex0 + ex1 + ex2 + ex3; float invZ = __frcp_rn(Z); float q[4] = { ex0 * invZ, ex1 * invZ, ex2 * invZ, ex3 * invZ }; float lZ = __logf(Z); float lsm[4] = { (out5[u][0] - mx) - lZ, (out5[u][1] - mx) - lZ, (out5[u][2] - mx) - lZ, (out5[u][3] - mx) - lZ }; float ratio = __expf(lsm[act] - mt[u].x); float H = -(q[0] * lsm[0] + q[1] * lsm[1] + q[2] * lsm[2] + q[3] * lsm[3]); float gate = ((advn > 0.0f && ratio > 1.2f) || (advn < 0.0f && ratio < 0.8f)) ? 0.0f : 1.0f; float pgw = -gate * advn * ratio * INV_MB; float dz[4]; #pragma unroll for (int m = 0; m < 4; ++m) { dz[m] = (pgw * ((m == act) ? 1.0f - q[m] : -q[m]) + (ENT_COEF * INV_MB) * q[m] * (H + lsm[m])) * valid[u]; } float dv = INV_MB * (out5[u][4] - mt[u].w) * valid[u]; float dh1 = w_dot_head(sm->w, lane, dz, dv); float dh2 = w_dot_head(sm->w, lane + 32, dz, dv); float dp1 = dh1 * (1.0f - h1[u] * h1[u]); float dp2 = dh2 * (1.0f - h2[u] * h2[u]); #pragma unroll for (int k = 0; k < 4; ++k) { a_w1[k] += dp1 * obs[u][k]; a_w1[4 + k] += dp2 * obs[u][k]; } a_b1[0] += dp1; a_b1[1] += dp2; #pragma unroll for (int m = 0; m < 4; ++m) { a_w2[m] += dz[m] * h1[u]; a_w2[4 + m] += dz[m] * h2[u]; } if (lane < 4) a_b2 += dz[lane]; a_w3[0] += dv * h1[u]; a_w3[1] += dv * h2[u]; if (lane == 0) a_b3 += dv; } } // commit per-lane accumulators -> shared-memory grad slots #pragma unroll for (int k = 0; k < 4; ++k) { atomicAdd(&sm->gacc[k * 64 + lane], a_w1[k]); atomicAdd(&sm->gacc[k * 64 + 32 + lane], a_w1[4 + k]); } atomicAdd(&sm->gacc[256 + lane], a_b1[0]); atomicAdd(&sm->gacc[256 + 32 + lane], a_b1[1]); #pragma unroll for (int m = 0; m < 4; ++m) { atomicAdd(&sm->gacc[320 + m * 64 + lane], a_w2[m]); atomicAdd(&sm->gacc[320 + m * 64 + 32 + lane], a_w2[4 + m]); } if (lane < 4) atomicAdd(&sm->gacc[576 + lane], a_b2); atomicAdd(&sm->gacc[580 + lane], a_w3[0]); atomicAdd(&sm->gacc[580 + 32 + lane], a_w3[1]); if (lane == 0) atomicAdd(&sm->gacc[644], a_b3); __syncthreads(); for (int k = tid; k < NPARAM; k += NT) { atomicAdd(&gslotG[slot0 + k], sm->gacc[k]); } grid.sync(); // ---- identical grad-clip + Adam in every block ------------- float part = 0.0f; for (int k = tid; k < NPARAM; k += NT) { float g = gslotG[slot0 + k]; gvals[k] = g; part += g * g; } #pragma unroll for (int off = 16; off >= 1; off >>= 1) part += __shfl_xor_sync(0xffffffffu, part, off); if (lane == 0) sm->red[wid] = part; __syncthreads(); float gtot = 0.0f; #pragma unroll for (int i = 0; i < NWARP; ++i) gtot += sm->red[i]; const float clip_coef = (gtot > 0.0f) ? fminf(1.0f, MAXN / (__fsqrt_rn(gtot) + 1e-6f)) : 1.0f; const int t_step = it * N_MB + midx + 1; const float bc1 = 1.0f - __powf(BETA1, (float)t_step); const float bc2 = 1.0f - __powf(BETA2, (float)t_step); for (int k = tid; k < NPARAM; k += NT) { float g = gvals[k] * clip_coef; float m1 = BETA1 * sm->am[k] + (1.0f - BETA1) * g; float v1 = BETA2 * sm->av[k] + (1.0f - BETA2) * g * g; sm->am[k] = m1; sm->av[k] = v1; float mhat = m1 / bc1; float vhat = v1 / bc2; sm->w[k] -= LRF * mhat / (__fsqrt_rn(vhat) + A_EPS); } __syncthreads(); } } } } // -------------------------------------------------------------------------- // two instantiations: NT=512 (preferred: fewest barrier arrivals) and NT=256 // (fallback if the 512-thread variant does not fit on the grading arch) // -------------------------------------------------------------------------- extern "C" __global__ void __launch_bounds__(512) ppo_k512(float4* o, float4* m, float* c, double* st, float* g, int* l, int* cn, int it, uint64_t sd) { ppo_body<512, 16>(o, m, c, st, g, l, cn, it, sd); } extern "C" __global__ void __launch_bounds__(256) ppo_k256(float4* o, float4* m, float* c, double* st, float* g, int* l, int* cn, int it, uint64_t sd) { ppo_body<256, 8>(o, m, c, st, g, l, cn, it, sd); } static int choose_nt() { static int cached = -1; if (cached > 0) return cached; int dev = 0; cudaGetDevice(&dev); int coop = 0; cudaDeviceGetAttribute(&coop, cudaDevAttrCooperativeLaunch, dev); TORCH_CHECK(coop, "device does not support cooperative launch"); int per_sm = 0; cudaError_t oerr = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per_sm, (const void*)ppo_k512, 512, 0); int nt = 256; if (oerr == cudaSuccess && per_sm > 0) nt = 512; else { cudaError_t o2 = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per_sm, (const void*)ppo_k256, 256, 0); TORCH_CHECK(o2 == cudaSuccess && per_sm > 0, "kernel does not fit on any SM"); } cached = nt; return nt; } int64_t query_nt() { return choose_nt(); } void run(torch::Tensor obsQ, torch::Tensor metaQ, torch::Tensor curve, torch::Tensor stats, torch::Tensor gslot, torch::Tensor listT, torch::Tensor cntT, int64_t iters, int64_t seed) { const int nt = choose_nt(); const void* func = (nt == 512) ? (const void*)ppo_k512 : (const void*)ppo_k256; int dev = 0; cudaGetDevice(&dev); cudaDeviceProp prop; cudaGetDeviceProperties(&prop, dev); int per_sm = 0; cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per_sm, func, nt, 0); int nb = per_sm * prop.multiProcessorCount; if (nb > 512) nb = 512; float4* obsP = (float4*)obsQ.data_ptr(); float4* metaP = (float4*)metaQ.data_ptr(); float* curveP = curve.data_ptr(); double* statsP = stats.data_ptr(); float* gslotP = gslot.data_ptr(); int* listP = (int*)listT.data_ptr(); int* cntP = (int*)cntT.data_ptr(); int it = (int)iters; uint64_t sd = (uint64_t)seed; void* args[] = { (void*)&obsP, (void*)&metaP, (void*)&curveP, (void*)&statsP, (void*)&gslotP, (void*)&listP, (void*)&cntP, (void*)&it, (void*)&sd }; cudaStream_t stream = at::cuda::getCurrentCUDAStream(); cudaError_t err = cudaLaunchCooperativeKernel(func, dim3(nb), dim3(nt), args, 0, stream); TORCH_CHECK(err == cudaSuccess, "cooperative launch failed: ", cudaGetErrorString(err)); cudaError_t lerr = cudaGetLastError(); TORCH_CHECK(lerr == cudaSuccess, "kernel launch error: ", cudaGetErrorString(lerr)); } """ def _build(): from torch.utils.cpp_extension import load_inline p = torch.cuda.get_device_properties(0) os.environ["TORCH_CUDA_ARCH_LIST"] = f"{p.major}.{p.minor}" return load_inline( name="ppo_megakernel_v11", cpp_sources=( "void run(torch::Tensor obsQ, torch::Tensor metaQ, torch::Tensor curve," " torch::Tensor stats, torch::Tensor gslot, torch::Tensor listT, torch::Tensor cntT," " int64_t iters, int64_t seed);" "int64_t query_nt();" ), cuda_sources=_CUDA_SRC, functions=["run", "query_nt"], extra_cuda_cflags=["-O3", "--restrict"], verbose=False, ) _EXT = None def _ext(): global _EXT if _EXT is None: _EXT = _build() return _EXT def train(total_env_steps: int, seed: int) -> list[float]: iters = max(1, int(total_env_steps) // (ROLLOUT * NUM_ENVS)) dev = torch.device("cuda:0") ext = _ext() obsQ = torch.empty(N_SAMPLES, 4, device=dev, dtype=torch.float32) metaQ = torch.empty(N_SAMPLES, 4, device=dev, dtype=torch.float32) curve = torch.zeros(iters, device=dev, dtype=torch.float32) stats = torch.zeros(2 * iters, device=dev, dtype=torch.float64) gslot = torch.zeros(iters * N_MINIBATCHES * NPARAM, device=dev, dtype=torch.float32) listT = torch.empty(N_MINIBATCHES * 40960, device=dev, dtype=torch.int32) cntT = torch.zeros(iters * N_MINIBATCHES, device=dev, dtype=torch.int32) ext.run(obsQ, metaQ, curve, stats, gslot, listT, cntT, iters, int(seed)) torch.cuda.synchronize() return (curve * (1.0 / NUM_ENVS)).cpu().tolist() if __name__ == "__main__": import time t0 = time.perf_counter() c = train(ROLLOUT * NUM_ENVS * 40, seed=0) torch.cuda.synchronize() dt = time.perf_counter() - t0 print(f"iters={len(c)} first={c[0]:.3f} last={c[-1]:.3f}") print(f"curve_head={[round(x,2) for x in c[:5]]} tail={[round(x,2) for x in c[-5:]]}") print(f"elapsed={dt:.3f}s sps={ROLLOUT * NUM_ENVS * len(c) / dt:,.0f}")