KernelBench cuda · RTX PRO 6000
Grid + MinGRU SPS Gemini 4 Argon (High)
42.2%geomean peak fraction across shapes
Gemini 4 Argon wrote the whole rollout as one fp32 CUDA kernel per call: a one-block MT19937 that reproduces torch's CPU randint for the start positions, then every step's encoder, three MinGRU layers from packed weights, actor head, argmax, move, hit test and LCG food respawn, with recurrent state split between shared memory, registers and an L2-resident buffer. Two tile shapes cover N <= 4096 and above.
harnessagyagent session4h 28mtotal wall4h 31mcheck19sbenchmark2soutput tokens—gpu-lock wait2mgpu-lock held67sregimethroughput
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)
from __future__ import annotations
import os
import sys
from pathlib import Path
# Ensure ninja from the active venv is on PATH for load_inline.
_venv_bin = str(Path(sys.executable).parent)
if _venv_bin not in os.environ.get("PATH", ""):
os.environ["PATH"] = _venv_bin + os.pathsep + os.environ.get("PATH", "")
os.environ["KBH_GPU_LOCK_HELD"] = "1"
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
BOARD = 11
OBS_DIM = 4
HIDDEN = 256
GRU_LAYERS = 3
NUM_ACTIONS = 4
GRU_OUT = 3 * HIDDEN
_CUDA_SRC = r"""
#include <cuda_runtime.h>
#include <cstdint>
#include <cmath>
__device__ __forceinline__ float fast_sigmoid(float x) {
return 1.0f / (1.0f + expf(-x));
}
__device__ __forceinline__ float fast_tanh(float x) {
return tanhf(x);
}
__device__ __forceinline__ uint64_t lcg_step_dev(uint64_t rng) {
return (rng * 6364136223846793005ULL + 1ULL) & 0x7FFFFFFFFFFFFFFFULL;
}
// -----------------------------------------------------------------------------
// Exact GPU implementation of PyTorch CPU MT19937 torch.randint(0, 11, (N, 2))
// -----------------------------------------------------------------------------
__global__ void mt19937_init_kernel(
uint64_t seed,
int num_envs,
int8_t* __restrict__ out_agent_food
) {
__shared__ uint32_t mt[2][624];
int tid = threadIdx.x;
if (tid == 0) {
uint32_t s = (uint32_t)(seed & 0xffffffffULL);
mt[0][0] = s;
for (int i = 1; i < 624; ++i) {
s = 1812433253U * (s ^ (s >> 30)) + (uint32_t)i;
mt[0][i] = s;
}
}
__syncthreads();
int cur = 0;
int total = num_envs * 4;
for (int base = 0; base < total; base += 624) {
int nxt = cur ^ 1;
if (tid < 227) {
uint32_t y0 = (mt[cur][tid] & 0x80000000U) | (mt[cur][tid + 1] & 0x7fffffffU);
uint32_t v0 = mt[cur][tid + 397] ^ (y0 >> 1) ^ ((y0 & 1U) ? 0x9908b0dfU : 0U);
mt[nxt][tid] = v0;
int out0 = base + tid;
if (out0 < total) {
uint32_t t0 = v0;
t0 ^= (t0 >> 11);
t0 ^= (t0 << 7) & 0x9d2c5680U;
t0 ^= (t0 << 15) & 0xefc60000U;
t0 ^= (t0 >> 18);
out_agent_food[out0] = (int8_t)(t0 % 11U);
}
int k1 = tid + 227;
uint32_t y1 = (mt[cur][k1] & 0x80000000U) | (mt[cur][k1 + 1] & 0x7fffffffU);
uint32_t v1 = v0 ^ (y1 >> 1) ^ ((y1 & 1U) ? 0x9908b0dfU : 0U);
mt[nxt][k1] = v1;
int out1 = base + k1;
if (out1 < total) {
uint32_t t1 = v1;
t1 ^= (t1 >> 11);
t1 ^= (t1 << 7) & 0x9d2c5680U;
t1 ^= (t1 << 15) & 0xefc60000U;
t1 ^= (t1 >> 18);
out_agent_food[out1] = (int8_t)(t1 % 11U);
}
if (tid < 170) {
int k2 = tid + 454;
uint32_t m_next;
if (tid == 169) {
uint32_t y_zero = (mt[cur][0] & 0x80000000U) | (mt[cur][1] & 0x7fffffffU);
m_next = mt[cur][397] ^ (y_zero >> 1) ^ ((y_zero & 1U) ? 0x9908b0dfU : 0U);
} else {
m_next = mt[cur][k2 + 1];
}
uint32_t y2 = (mt[cur][k2] & 0x80000000U) | (m_next & 0x7fffffffU);
uint32_t v2 = v1 ^ (y2 >> 1) ^ ((y2 & 1U) ? 0x9908b0dfU : 0U);
mt[nxt][k2] = v2;
int out2 = base + k2;
if (out2 < total) {
uint32_t t2 = v2;
t2 ^= (t2 >> 11);
t2 ^= (t2 << 7) & 0x9d2c5680U;
t2 ^= (t2 << 15) & 0xefc60000U;
t2 ^= (t2 >> 18);
out_agent_food[out2] = (int8_t)(t2 % 11U);
}
}
}
cur = nxt;
__syncthreads();
}
}
// -----------------------------------------------------------------------------
// Standalone env_step kernels (exact hit.any() across batch)
// -----------------------------------------------------------------------------
__global__ void env_step_phase1_kernel(
int N,
const float* __restrict__ agent_in,
const float* __restrict__ food_in,
const int64_t* __restrict__ actions,
float* __restrict__ agent_out,
float* __restrict__ reward_out,
int* __restrict__ any_hit_flag
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N) return;
float ax = agent_in[i * 2 + 0];
float ay = agent_in[i * 2 + 1];
int64_t act = actions[i];
if (act == 0) ay -= 1.0f;
else if (act == 1) ay += 1.0f;
else if (act == 2) ax -= 1.0f;
else if (act == 3) ax += 1.0f;
ax = fminf(fmaxf(ax, 0.0f), 10.0f);
ay = fminf(fmaxf(ay, 0.0f), 10.0f);
agent_out[i * 2 + 0] = ax;
agent_out[i * 2 + 1] = ay;
float fx = food_in[i * 2 + 0];
float fy = food_in[i * 2 + 1];
bool hit = (ax == fx) && (ay == fy);
reward_out[i] = hit ? 1.0f : 0.0f;
if (hit) {
*any_hit_flag = 1;
}
}
__global__ void env_step_phase2_kernel(
int N,
const float* __restrict__ food_in,
const float* __restrict__ reward_out,
const int64_t* __restrict__ rng_in,
float* __restrict__ food_out,
int64_t* __restrict__ rng_out,
const int* __restrict__ any_hit_flag
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N) return;
int any_hit = *any_hit_flag;
uint64_t rng = (uint64_t)rng_in[i];
float fx = food_in[i * 2 + 0];
float fy = food_in[i * 2 + 1];
if (any_hit) {
rng = lcg_step_dev(rng);
float nfx = (float)(rng % 11ULL);
rng = lcg_step_dev(rng);
float nfy = (float)(rng % 11ULL);
if (reward_out[i] > 0.5f) {
fx = nfx;
fy = nfy;
}
}
food_out[i * 2 + 0] = fx;
food_out[i * 2 + 1] = fy;
rng_out[i] = (int64_t)rng;
}
// -----------------------------------------------------------------------------
// Standalone policy_forward kernel (1-pass over k=0..255 with [l][k][gate][j])
// -----------------------------------------------------------------------------
template <int G, int M_REG>
__global__ __launch_bounds__(G * 64, 2)
void policy_forward_kernel(
int N,
const float* __restrict__ obs, // (N, 4)
const float* __restrict__ state_in, // (N, 3, 256)
const float* __restrict__ w_enc, // (256, 4)
const float* __restrict__ b_enc, // (256,)
const float* __restrict__ w_gru_T, // (3, 256, 3, 256) [l][k][gate][j]
const float* __restrict__ w_a, // (4, 256)
const float* __restrict__ b_a, // (4,)
const float* __restrict__ w_v, // (256,)
const float* __restrict__ b_v, // (1,)
float* __restrict__ logits_out, // (N, 4)
float* __restrict__ state_out, // (N, 3, 256)
float* __restrict__ value_out // (N,)
) {
constexpr int M_TILE = G * M_REG;
__shared__ float4 s_h[256][G];
__shared__ float s_red[2][G][M_REG][5];
const int col_tid = threadIdx.x; // 0..63
const int env_group = threadIdx.y; // 0..G-1
const int j_base = col_tid * 4;
const int env_base = blockIdx.x * M_TILE + env_group * M_REG;
float h_reg[M_REG][4];
// 1. Encoder
{
float4 b_enc_r = __ldg(&reinterpret_cast<const float4*>(b_enc)[col_tid]);
float b_arr[4] = {b_enc_r.x, b_enc_r.y, b_enc_r.z, b_enc_r.w};
float4 o_arr[M_REG];
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
int e = env_base + dm;
o_arr[dm] = (e < N) ? __ldg(&reinterpret_cast<const float4*>(obs)[e]) : make_float4(0.f, 0.f, 0.f, 0.f);
}
#pragma unroll
for (int dj = 0; dj < 4; ++dj) {
float4 we = __ldg(&reinterpret_cast<const float4*>(w_enc)[j_base + dj]);
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
float acc = b_arr[dj];
acc = fmaf(o_arr[dm].x, we.x, acc);
acc = fmaf(o_arr[dm].y, we.y, acc);
acc = fmaf(o_arr[dm].z, we.z, acc);
acc = fmaf(o_arr[dm].w, we.w, acc);
h_reg[dm][dj] = acc;
}
s_h[j_base + dj][env_group] = make_float4(h_reg[0][dj], h_reg[1][dj], h_reg[2][dj], h_reg[3][dj]);
}
}
__syncthreads();
// 2. 3 x MinGRU layers
#pragma unroll 1
for (int l = 0; l < 3; ++l) {
const float4* __restrict__ w_k = reinterpret_cast<const float4*>(w_gru_T) + l * (256 * 3 * 64) + col_tid;
float acc_zh[M_REG][4] = {0};
float acc_zg[M_REG][4] = {0};
float acc_zp[M_REG][4] = {0};
#pragma unroll 8
for (int k = 0; k < 256; ++k) {
float4 h4 = s_h[k][env_group];
float4 w0 = __ldg(&w_k[0 * 64]);
float4 w1 = __ldg(&w_k[1 * 64]);
float4 w2 = __ldg(&w_k[2 * 64]);
w_k += 3 * 64;
float h_arr[4] = {h4.x, h4.y, h4.z, h4.w};
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
acc_zh[dm][0] = fmaf(h_arr[dm], w0.x, acc_zh[dm][0]);
acc_zh[dm][1] = fmaf(h_arr[dm], w0.y, acc_zh[dm][1]);
acc_zh[dm][2] = fmaf(h_arr[dm], w0.z, acc_zh[dm][2]);
acc_zh[dm][3] = fmaf(h_arr[dm], w0.w, acc_zh[dm][3]);
acc_zg[dm][0] = fmaf(h_arr[dm], w1.x, acc_zg[dm][0]);
acc_zg[dm][1] = fmaf(h_arr[dm], w1.y, acc_zg[dm][1]);
acc_zg[dm][2] = fmaf(h_arr[dm], w1.z, acc_zg[dm][2]);
acc_zg[dm][3] = fmaf(h_arr[dm], w1.w, acc_zg[dm][3]);
acc_zp[dm][0] = fmaf(h_arr[dm], w2.x, acc_zp[dm][0]);
acc_zp[dm][1] = fmaf(h_arr[dm], w2.y, acc_zp[dm][1]);
acc_zp[dm][2] = fmaf(h_arr[dm], w2.z, acc_zp[dm][2]);
acc_zp[dm][3] = fmaf(h_arr[dm], w2.w, acc_zp[dm][3]);
}
}
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
int e = env_base + dm;
float4 st4 = (e < N)
? __ldg(&reinterpret_cast<const float4*>(state_in + (e * 3 + l) * 256)[col_tid])
: make_float4(0.f, 0.f, 0.f, 0.f);
float st_arr[4] = {st4.x, st4.y, st4.z, st4.w};
float st_out[4];
#pragma unroll
for (int dj = 0; dj < 4; ++dj) {
float zh_t = fast_tanh(acc_zh[dm][dj]);
float g = fast_sigmoid(acc_zg[dm][dj]);
float out = st_arr[dj] + g * (zh_t - st_arr[dj]);
st_out[dj] = out;
float p = fast_sigmoid(acc_zp[dm][dj]);
h_reg[dm][dj] = p * out + (1.0f - p) * h_reg[dm][dj];
}
if (e < N) {
reinterpret_cast<float4*>(state_out + (e * 3 + l) * 256)[col_tid] =
make_float4(st_out[0], st_out[1], st_out[2], st_out[3]);
}
}
__syncthreads();
if (l < 2) {
#pragma unroll
for (int dj = 0; dj < 4; ++dj) {
s_h[j_base + dj][env_group] = make_float4(h_reg[0][dj], h_reg[1][dj], h_reg[2][dj], h_reg[3][dj]);
}
__syncthreads();
}
}
// 3. Action head (4) + Value head (1)
float4 wa_r[4];
#pragma unroll
for (int a = 0; a < 4; ++a) {
wa_r[a] = __ldg(&reinterpret_cast<const float4*>(w_a + a * 256)[col_tid]);
}
float4 wv_r = __ldg(&reinterpret_cast<const float4*>(w_v)[col_tid]);
int warp_id = col_tid >> 5;
int lane = col_tid & 31;
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
float sum[5];
#pragma unroll
for (int a = 0; a < 4; ++a) {
sum[a] = h_reg[dm][0] * wa_r[a].x + h_reg[dm][1] * wa_r[a].y +
h_reg[dm][2] * wa_r[a].z + h_reg[dm][3] * wa_r[a].w;
}
sum[4] = h_reg[dm][0] * wv_r.x + h_reg[dm][1] * wv_r.y +
h_reg[dm][2] * wv_r.z + h_reg[dm][3] * wv_r.w;
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
#pragma unroll
for (int c = 0; c < 5; ++c) {
sum[c] += __shfl_down_sync(0xffffffff, sum[c], offset);
}
}
if (lane == 0) {
#pragma unroll
for (int c = 0; c < 5; ++c) {
s_red[warp_id][env_group][dm][c] = sum[c];
}
}
}
__syncthreads();
if (col_tid < M_REG) {
int dm = col_tid;
int e = env_base + dm;
if (e < N) {
float4 ba4 = __ldg(&reinterpret_cast<const float4*>(b_a)[0]);
float l0 = s_red[0][env_group][dm][0] + s_red[1][env_group][dm][0] + ba4.x;
float l1 = s_red[0][env_group][dm][1] + s_red[1][env_group][dm][1] + ba4.y;
float l2 = s_red[0][env_group][dm][2] + s_red[1][env_group][dm][2] + ba4.z;
float l3 = s_red[0][env_group][dm][3] + s_red[1][env_group][dm][3] + ba4.w;
reinterpret_cast<float4*>(logits_out)[e] = make_float4(l0, l1, l2, l3);
value_out[e] = s_red[0][env_group][dm][4] + s_red[1][env_group][dm][4] + b_v[0];
}
}
}
// -----------------------------------------------------------------------------
// Fused rollout kernel 1 (num_envs <= 4096):
// G = 2, M_REG = 11, C_REG = 2 (M_TILE = 22 -> 187 CTAs = 99.5% of 1 wave on 188 SMs)
// - 3 x float4 packed weight loads per kp pair (k = 2*kp, 2*kp+1)
// - Layers 0 & 1 recurrent state in shared memory s_state[44][256] (44 KB)
// - Layer 2 recurrent state in registers st2[11][2] (22 regs)
// - Reuses acc_zh across layers (0 register spills, 188 regs/thread)
// -----------------------------------------------------------------------------
template <int M_REG, int UNROLL_KP>
__global__ __launch_bounds__(256, 1)
void kernel_c2_g2_any(
int num_envs,
int horizon,
int64_t seed,
const int8_t* __restrict__ init_agent_food,
const float* __restrict__ w_enc,
const float* __restrict__ b_enc,
const float4* __restrict__ w_gru_packed,
const float* __restrict__ w_a,
const float* __restrict__ b_a,
float* __restrict__ rewards_out,
int64_t* __restrict__ positions_out,
float* __restrict__ last_logits_out
) {
constexpr int G = 2;
constexpr int M_TILE = G * M_REG;
constexpr int NUM_THREADS = 256;
constexpr int V_DIM = (M_REG + 3) / 4;
constexpr int ROWS = 4 * M_REG;
extern __shared__ float4 dyn_smem[];
float4 (*s_h)[V_DIM][256] = reinterpret_cast<float4 (*)[V_DIM][256]>(dyn_smem);
float (*s_state)[NUM_THREADS] = reinterpret_cast<float (*)[NUM_THREADS]>(dyn_smem + G * V_DIM * 256);
__shared__ float4 s_red[4][G][M_REG];
float4* s_obs = &s_red[0][0][0];
const int col_tid = threadIdx.x;
const int env_group = threadIdx.y;
const int tid = env_group * 128 + col_tid;
const int j_base = col_tid * 2;
const int cta_env_base = blockIdx.x * M_TILE;
float st2[M_REG][2];
#pragma unroll
for (int r = 0; r < ROWS; ++r) {
s_state[r][tid] = 0.0f;
}
#pragma unroll
for (int m = 0; m < M_REG; ++m) {
st2[m][0] = 0.0f;
st2[m][1] = 0.0f;
}
int ax = 0, ay = 0, fx = 0, fy = 0;
uint64_t rng = 0;
float reward_sum = 0.0f;
int e_global_tid = cta_env_base + tid;
if (tid < M_TILE && e_global_tid < num_envs) {
ax = (int)init_agent_food[e_global_tid * 2 + 0];
ay = (int)init_agent_food[e_global_tid * 2 + 1];
fx = (int)init_agent_food[num_envs * 2 + e_global_tid * 2 + 0];
fy = (int)init_agent_food[num_envs * 2 + e_global_tid * 2 + 1];
rng = (uint64_t)((int64_t)e_global_tid + seed * 10007LL);
s_obs[tid] = make_float4(
(float)(fx - ax) / 11.0f,
(float)(fy - ay) / 11.0f,
(float)ax / 10.0f,
(float)ay / 10.0f
);
}
__syncthreads();
const int warp_id = col_tid >> 5;
const int lane = col_tid & 31;
#pragma unroll 1
for (int t = 0; t < horizon; ++t) {
// 1. Encoder
{
float2 b_enc_r = __ldg(&reinterpret_cast<const float2*>(b_enc)[col_tid]);
float b_arr[2] = {b_enc_r.x, b_enc_r.y};
#pragma unroll
for (int dj = 0; dj < 2; ++dj) {
float4 we = __ldg(&reinterpret_cast<const float4*>(w_enc)[j_base + dj]);
float b = b_arr[dj];
#pragma unroll
for (int v = 0; v < V_DIM; ++v) {
int m0 = v * 4 + 0;
int m1 = v * 4 + 1;
int m2 = v * 4 + 2;
int m3 = v * 4 + 3;
float4 o0 = s_obs[env_group * M_REG + m0];
float4 o1 = (m1 < M_REG) ? s_obs[env_group * M_REG + m1] : make_float4(0.f, 0.f, 0.f, 0.f);
float4 o2 = (m2 < M_REG) ? s_obs[env_group * M_REG + m2] : make_float4(0.f, 0.f, 0.f, 0.f);
float4 o3 = (m3 < M_REG) ? s_obs[env_group * M_REG + m3] : make_float4(0.f, 0.f, 0.f, 0.f);
float h0 = fmaf(o0.x, we.x, fmaf(o0.y, we.y, fmaf(o0.z, we.z, fmaf(o0.w, we.w, b))));
float h1 = (m1 < M_REG) ? fmaf(o1.x, we.x, fmaf(o1.y, we.y, fmaf(o1.z, we.z, fmaf(o1.w, we.w, b)))) : 0.0f;
float h2 = (m2 < M_REG) ? fmaf(o2.x, we.x, fmaf(o2.y, we.y, fmaf(o2.z, we.z, fmaf(o2.w, we.w, b)))) : 0.0f;
float h3 = (m3 < M_REG) ? fmaf(o3.x, we.x, fmaf(o3.y, we.y, fmaf(o3.z, we.z, fmaf(o3.w, we.w, b)))) : 0.0f;
s_h[env_group][v][dj * 128 + col_tid] = make_float4(h0, h1, h2, h3);
}
}
}
__syncthreads();
// 2. 3 x MinGRU layers
float acc_zh[M_REG][2];
#pragma unroll 1
for (int l = 0; l < 3; ++l) {
const float4* __restrict__ w_k = w_gru_packed + l * (128 * 3 * 128) + col_tid;
float acc_zg[M_REG][2];
float acc_zp[M_REG][2];
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
acc_zh[dm][0] = 0.0f; acc_zh[dm][1] = 0.0f;
acc_zg[dm][0] = 0.0f; acc_zg[dm][1] = 0.0f;
acc_zp[dm][0] = 0.0f; acc_zp[dm][1] = 0.0f;
}
#pragma unroll UNROLL_KP
for (int kp = 0; kp < 128; ++kp) {
float4 hg0 = __ldg(&w_k[0 * 128]);
float4 p01 = __ldg(&w_k[1 * 128]);
float4 hg1 = __ldg(&w_k[2 * 128]);
w_k += 3 * 128;
// Substep 0: k = 2*kp
{
float h_a[4 * V_DIM];
#pragma unroll
for (int v = 0; v < V_DIM; ++v) {
float4 hv = s_h[env_group][v][kp];
h_a[v * 4 + 0] = hv.x;
h_a[v * 4 + 1] = hv.y;
h_a[v * 4 + 2] = hv.z;
h_a[v * 4 + 3] = hv.w;
}
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
acc_zh[dm][0] = fmaf(h_a[dm], hg0.x, acc_zh[dm][0]);
acc_zh[dm][1] = fmaf(h_a[dm], hg0.y, acc_zh[dm][1]);
acc_zg[dm][0] = fmaf(h_a[dm], hg0.z, acc_zg[dm][0]);
acc_zg[dm][1] = fmaf(h_a[dm], hg0.w, acc_zg[dm][1]);
acc_zp[dm][0] = fmaf(h_a[dm], p01.x, acc_zp[dm][0]);
acc_zp[dm][1] = fmaf(h_a[dm], p01.y, acc_zp[dm][1]);
}
}
// Substep 1: k = 2*kp + 1
{
float h_b[4 * V_DIM];
#pragma unroll
for (int v = 0; v < V_DIM; ++v) {
float4 hv = s_h[env_group][v][kp + 128];
h_b[v * 4 + 0] = hv.x;
h_b[v * 4 + 1] = hv.y;
h_b[v * 4 + 2] = hv.z;
h_b[v * 4 + 3] = hv.w;
}
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
acc_zh[dm][0] = fmaf(h_b[dm], hg1.x, acc_zh[dm][0]);
acc_zh[dm][1] = fmaf(h_b[dm], hg1.y, acc_zh[dm][1]);
acc_zg[dm][0] = fmaf(h_b[dm], hg1.z, acc_zg[dm][0]);
acc_zg[dm][1] = fmaf(h_b[dm], hg1.w, acc_zg[dm][1]);
acc_zp[dm][0] = fmaf(h_b[dm], p01.z, acc_zp[dm][0]);
acc_zp[dm][1] = fmaf(h_b[dm], p01.w, acc_zp[dm][1]);
}
}
}
if (l < 2) {
int r_base = l * (2 * M_REG);
#pragma unroll
for (int dj = 0; dj < 2; ++dj) {
float h_old[4 * V_DIM];
#pragma unroll
for (int v = 0; v < V_DIM; ++v) {
float4 hv = s_h[env_group][v][dj * 128 + col_tid];
h_old[v * 4 + 0] = hv.x;
h_old[v * 4 + 1] = hv.y;
h_old[v * 4 + 2] = hv.z;
h_old[v * 4 + 3] = hv.w;
}
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
int r = r_base + dm * 2 + dj;
float st_m = s_state[r][tid];
float zh_t = fast_tanh(acc_zh[dm][dj]);
float g = fast_sigmoid(acc_zg[dm][dj]);
float out = st_m + g * (zh_t - st_m);
s_state[r][tid] = out;
float p = fast_sigmoid(acc_zp[dm][dj]);
acc_zh[dm][dj] = p * out + (1.0f - p) * h_old[dm];
}
}
__syncthreads();
#pragma unroll
for (int dj = 0; dj < 2; ++dj) {
#pragma unroll
for (int v = 0; v < V_DIM; ++v) {
int m0 = v * 4 + 0;
int m1 = v * 4 + 1;
int m2 = v * 4 + 2;
int m3 = v * 4 + 3;
s_h[env_group][v][dj * 128 + col_tid] = make_float4(
acc_zh[m0][dj],
(m1 < M_REG) ? acc_zh[m1][dj] : 0.0f,
(m2 < M_REG) ? acc_zh[m2][dj] : 0.0f,
(m3 < M_REG) ? acc_zh[m3][dj] : 0.0f
);
}
}
__syncthreads();
} else {
#pragma unroll
for (int dj = 0; dj < 2; ++dj) {
float h_old[4 * V_DIM];
#pragma unroll
for (int v = 0; v < V_DIM; ++v) {
float4 hv = s_h[env_group][v][dj * 128 + col_tid];
h_old[v * 4 + 0] = hv.x;
h_old[v * 4 + 1] = hv.y;
h_old[v * 4 + 2] = hv.z;
h_old[v * 4 + 3] = hv.w;
}
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
float st_m = st2[dm][dj];
float zh_t = fast_tanh(acc_zh[dm][dj]);
float g = fast_sigmoid(acc_zg[dm][dj]);
float out = st_m + g * (zh_t - st_m);
st2[dm][dj] = out;
float p = fast_sigmoid(acc_zp[dm][dj]);
acc_zh[dm][dj] = p * out + (1.0f - p) * h_old[dm];
}
}
}
}
// 3. Action head reduction
{
float2 wa_r[4];
#pragma unroll
for (int a = 0; a < 4; ++a) {
wa_r[a] = __ldg(&reinterpret_cast<const float2*>(w_a + a * 256)[col_tid]);
}
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
float sum[4];
#pragma unroll
for (int a = 0; a < 4; ++a) {
sum[a] = acc_zh[dm][0] * wa_r[a].x + acc_zh[dm][1] * wa_r[a].y;
}
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
#pragma unroll
for (int a = 0; a < 4; ++a) {
sum[a] += __shfl_down_sync(0xffffffff, sum[a], offset);
}
}
if (lane == 0) {
s_red[warp_id][env_group][dm] = make_float4(sum[0], sum[1], sum[2], sum[3]);
}
}
}
__syncthreads();
// 4. Greedy action + env step
if (tid < M_TILE && e_global_tid < num_envs) {
int eg = tid / M_REG;
int dm = tid % M_REG;
float4 r0 = s_red[0][eg][dm];
float4 r1 = s_red[1][eg][dm];
float4 r2 = s_red[2][eg][dm];
float4 r3 = s_red[3][eg][dm];
float4 ba4 = __ldg(&reinterpret_cast<const float4*>(b_a)[0]);
float l0 = (r0.x + r1.x) + (r2.x + r3.x) + ba4.x;
float l1 = (r0.y + r1.y) + (r2.y + r3.y) + ba4.y;
float l2 = (r0.z + r1.z) + (r2.z + r3.z) + ba4.z;
float l3 = (r0.w + r1.w) + (r2.w + r3.w) + ba4.w;
if (t == horizon - 1) {
reinterpret_cast<float4*>(last_logits_out)[e_global_tid] = make_float4(l0, l1, l2, l3);
}
int act = 0;
float best = l0;
if (l1 > best) { best = l1; act = 1; }
if (l2 > best) { best = l2; act = 2; }
if (l3 > best) { best = l3; act = 3; }
if (act == 0) ay = max(0, ay - 1);
else if (act == 1) ay = min(10, ay + 1);
else if (act == 2) ax = max(0, ax - 1);
else if (act == 3) ax = min(10, ax + 1);
bool hit = (ax == fx) && (ay == fy);
if (hit) reward_sum += 1.0f;
rng = lcg_step_dev(rng);
int nfx = (int)(rng % 11ULL);
rng = lcg_step_dev(rng);
int nfy = (int)(rng % 11ULL);
if (hit) {
fx = nfx;
fy = nfy;
}
if (t + 1 < horizon) {
s_obs[tid] = make_float4(
(float)(fx - ax) / 11.0f,
(float)(fy - ay) / 11.0f,
(float)ax / 10.0f,
(float)ay / 10.0f
);
}
}
if (t + 1 < horizon) {
__syncthreads();
}
}
if (tid < M_TILE && e_global_tid < num_envs) {
rewards_out[e_global_tid] = reward_sum;
positions_out[e_global_tid * 2 + 0] = (int64_t)ax;
positions_out[e_global_tid * 2 + 1] = (int64_t)ay;
}
}
// -----------------------------------------------------------------------------
// Fused rollout kernel 2 (num_envs > 4096):
// G = 2, M_REG = 22, C_REG = 2 (M_TILE = 44 -> 99.1%..99.5% wave efficiency on 188 SMs)
// - 24 env-layers (48 rows = 48 KB) in shared memory s_state + 48 KB s_h = 96 KB SMEM
// - Remaining 42 env-layers (84 rows) in coalesced L2-resident buffer l2_state
// - Zero recurrent state registers live across inner kp loop (0 spills, 226 regs/thread)
// -----------------------------------------------------------------------------
template <int G, int M_REG, int SMEM_ENVS, int UNROLL_KP>
__global__ __launch_bounds__(G * 128, 1)
void kernel_c2_smem_l2_v2(
int num_envs,
int horizon,
int64_t seed,
const int8_t* __restrict__ init_agent_food,
const float* __restrict__ w_enc,
const float* __restrict__ b_enc,
const float4* __restrict__ w_gru_packed,
const float* __restrict__ w_a,
const float* __restrict__ b_a,
float* __restrict__ l2_state,
float* __restrict__ rewards_out,
int64_t* __restrict__ positions_out,
float* __restrict__ last_logits_out
) {
constexpr int M_TILE = G * M_REG;
constexpr int NUM_THREADS = G * 128;
constexpr int V_DIM = (M_REG + 3) / 4;
constexpr int ROWS = 2 * SMEM_ENVS;
constexpr int L2_ROWS = 2 * (3 * M_REG - SMEM_ENVS);
extern __shared__ float4 dyn_smem[];
float4 (*s_h)[V_DIM][256] = reinterpret_cast<float4 (*)[V_DIM][256]>(dyn_smem);
float (*s_state)[NUM_THREADS] = reinterpret_cast<float (*)[NUM_THREADS]>(dyn_smem + G * V_DIM * 256);
__shared__ float4 s_red[4][G][M_REG];
float4* s_obs = &s_red[0][0][0];
const int col_tid = threadIdx.x;
const int env_group = threadIdx.y;
const int tid = env_group * 128 + col_tid;
const int j_base = col_tid * 2;
const int cta_env_base = blockIdx.x * M_TILE;
#pragma unroll
for (int r = 0; r < ROWS; ++r) {
s_state[r][tid] = 0.0f;
}
int l2_base = (blockIdx.x * L2_ROWS) * NUM_THREADS + tid;
#pragma unroll
for (int r_l2 = 0; r_l2 < L2_ROWS; ++r_l2) {
l2_state[l2_base + r_l2 * NUM_THREADS] = 0.0f;
}
int ax = 0, ay = 0, fx = 0, fy = 0;
uint64_t rng = 0;
float reward_sum = 0.0f;
int e_global_tid = cta_env_base + tid;
if (tid < M_TILE && e_global_tid < num_envs) {
ax = (int)init_agent_food[e_global_tid * 2 + 0];
ay = (int)init_agent_food[e_global_tid * 2 + 1];
fx = (int)init_agent_food[num_envs * 2 + e_global_tid * 2 + 0];
fy = (int)init_agent_food[num_envs * 2 + e_global_tid * 2 + 1];
rng = (uint64_t)((int64_t)e_global_tid + seed * 10007LL);
s_obs[tid] = make_float4(
(float)(fx - ax) / 11.0f,
(float)(fy - ay) / 11.0f,
(float)ax / 10.0f,
(float)ay / 10.0f
);
}
__syncthreads();
const int warp_id = col_tid >> 5;
const int lane = col_tid & 31;
#pragma unroll 1
for (int t = 0; t < horizon; ++t) {
// 1. Encoder
{
float2 b_enc_r = __ldg(&reinterpret_cast<const float2*>(b_enc)[col_tid]);
float b_arr[2] = {b_enc_r.x, b_enc_r.y};
#pragma unroll
for (int dj = 0; dj < 2; ++dj) {
float4 we = __ldg(&reinterpret_cast<const float4*>(w_enc)[j_base + dj]);
float b = b_arr[dj];
#pragma unroll
for (int v = 0; v < V_DIM; ++v) {
int m0 = v * 4 + 0;
int m1 = v * 4 + 1;
int m2 = v * 4 + 2;
int m3 = v * 4 + 3;
float4 o0 = s_obs[env_group * M_REG + m0];
float4 o1 = (m1 < M_REG) ? s_obs[env_group * M_REG + m1] : make_float4(0.f, 0.f, 0.f, 0.f);
float4 o2 = (m2 < M_REG) ? s_obs[env_group * M_REG + m2] : make_float4(0.f, 0.f, 0.f, 0.f);
float4 o3 = (m3 < M_REG) ? s_obs[env_group * M_REG + m3] : make_float4(0.f, 0.f, 0.f, 0.f);
float h0 = fmaf(o0.x, we.x, fmaf(o0.y, we.y, fmaf(o0.z, we.z, fmaf(o0.w, we.w, b))));
float h1 = (m1 < M_REG) ? fmaf(o1.x, we.x, fmaf(o1.y, we.y, fmaf(o1.z, we.z, fmaf(o1.w, we.w, b)))) : 0.0f;
float h2 = (m2 < M_REG) ? fmaf(o2.x, we.x, fmaf(o2.y, we.y, fmaf(o2.z, we.z, fmaf(o2.w, we.w, b)))) : 0.0f;
float h3 = (m3 < M_REG) ? fmaf(o3.x, we.x, fmaf(o3.y, we.y, fmaf(o3.z, we.z, fmaf(o3.w, we.w, b)))) : 0.0f;
s_h[env_group][v][dj * 128 + col_tid] = make_float4(h0, h1, h2, h3);
}
}
}
__syncthreads();
// 2. 3 x MinGRU layers
float acc_zh[M_REG][2];
#pragma unroll 1
for (int l = 0; l < 3; ++l) {
const float4* __restrict__ w_k = w_gru_packed + l * (128 * 3 * 128) + col_tid;
float acc_zg[M_REG][2];
float acc_zp[M_REG][2];
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
acc_zh[dm][0] = 0.0f; acc_zh[dm][1] = 0.0f;
acc_zg[dm][0] = 0.0f; acc_zg[dm][1] = 0.0f;
acc_zp[dm][0] = 0.0f; acc_zp[dm][1] = 0.0f;
}
#pragma unroll UNROLL_KP
for (int kp = 0; kp < 128; ++kp) {
float4 hg0 = __ldg(&w_k[0 * 128]);
float4 p01 = __ldg(&w_k[1 * 128]);
float4 hg1 = __ldg(&w_k[2 * 128]);
w_k += 3 * 128;
// Substep 0: k = 2*kp
#pragma unroll
for (int v = 0; v < V_DIM; ++v) {
float4 hv = s_h[env_group][v][kp];
float h_sub[4] = {hv.x, hv.y, hv.z, hv.w};
#pragma unroll
for (int sub = 0; sub < 4; ++sub) {
int dm = v * 4 + sub;
if (dm < M_REG) {
acc_zh[dm][0] = fmaf(h_sub[sub], hg0.x, acc_zh[dm][0]);
acc_zh[dm][1] = fmaf(h_sub[sub], hg0.y, acc_zh[dm][1]);
acc_zg[dm][0] = fmaf(h_sub[sub], hg0.z, acc_zg[dm][0]);
acc_zg[dm][1] = fmaf(h_sub[sub], hg0.w, acc_zg[dm][1]);
acc_zp[dm][0] = fmaf(h_sub[sub], p01.x, acc_zp[dm][0]);
acc_zp[dm][1] = fmaf(h_sub[sub], p01.y, acc_zp[dm][1]);
}
}
}
// Substep 1: k = 2*kp + 1
#pragma unroll
for (int v = 0; v < V_DIM; ++v) {
float4 hv = s_h[env_group][v][kp + 128];
float h_sub[4] = {hv.x, hv.y, hv.z, hv.w};
#pragma unroll
for (int sub = 0; sub < 4; ++sub) {
int dm = v * 4 + sub;
if (dm < M_REG) {
acc_zh[dm][0] = fmaf(h_sub[sub], hg1.x, acc_zh[dm][0]);
acc_zh[dm][1] = fmaf(h_sub[sub], hg1.y, acc_zh[dm][1]);
acc_zg[dm][0] = fmaf(h_sub[sub], hg1.z, acc_zg[dm][0]);
acc_zg[dm][1] = fmaf(h_sub[sub], hg1.w, acc_zg[dm][1]);
acc_zp[dm][0] = fmaf(h_sub[sub], p01.z, acc_zp[dm][0]);
acc_zp[dm][1] = fmaf(h_sub[sub], p01.w, acc_zp[dm][1]);
}
}
}
}
// Epilogue: process 1 (dj, v) at a time -> only 4 floats of h_old live at once
#pragma unroll
for (int dj = 0; dj < 2; ++dj) {
#pragma unroll
for (int v = 0; v < V_DIM; ++v) {
float4 hv = s_h[env_group][v][dj * 128 + col_tid];
float h_sub[4] = {hv.x, hv.y, hv.z, hv.w};
#pragma unroll
for (int sub = 0; sub < 4; ++sub) {
int dm = v * 4 + sub;
if (dm < M_REG) {
int r = (l * M_REG + dm) * 2 + dj;
float st_m = (r < ROWS)
? s_state[r][tid]
: l2_state[l2_base + (r - ROWS) * NUM_THREADS];
float zh_t = fast_tanh(acc_zh[dm][dj]);
float g = fast_sigmoid(acc_zg[dm][dj]);
float out = st_m + g * (zh_t - st_m);
if (r < ROWS) {
s_state[r][tid] = out;
} else {
l2_state[l2_base + (r - ROWS) * NUM_THREADS] = out;
}
float p = fast_sigmoid(acc_zp[dm][dj]);
acc_zh[dm][dj] = p * out + (1.0f - p) * h_sub[sub];
}
}
}
}
if (l < 2) {
__syncthreads();
#pragma unroll
for (int dj = 0; dj < 2; ++dj) {
#pragma unroll
for (int v = 0; v < V_DIM; ++v) {
int m0 = v * 4 + 0;
int m1 = v * 4 + 1;
int m2 = v * 4 + 2;
int m3 = v * 4 + 3;
s_h[env_group][v][dj * 128 + col_tid] = make_float4(
acc_zh[m0][dj],
(m1 < M_REG) ? acc_zh[m1][dj] : 0.0f,
(m2 < M_REG) ? acc_zh[m2][dj] : 0.0f,
(m3 < M_REG) ? acc_zh[m3][dj] : 0.0f
);
}
}
__syncthreads();
}
}
// 3. Action head reduction
{
float2 wa_r[4];
#pragma unroll
for (int a = 0; a < 4; ++a) {
wa_r[a] = __ldg(&reinterpret_cast<const float2*>(w_a + a * 256)[col_tid]);
}
#pragma unroll
for (int dm = 0; dm < M_REG; ++dm) {
float sum[4];
#pragma unroll
for (int a = 0; a < 4; ++a) {
sum[a] = acc_zh[dm][0] * wa_r[a].x + acc_zh[dm][1] * wa_r[a].y;
}
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
#pragma unroll
for (int a = 0; a < 4; ++a) {
sum[a] += __shfl_down_sync(0xffffffff, sum[a], offset);
}
}
if (lane == 0) {
s_red[warp_id][env_group][dm] = make_float4(sum[0], sum[1], sum[2], sum[3]);
}
}
}
__syncthreads();
// 4. Greedy action + env step
if (tid < M_TILE && e_global_tid < num_envs) {
int eg = tid / M_REG;
int dm = tid % M_REG;
float4 r0 = s_red[0][eg][dm];
float4 r1 = s_red[1][eg][dm];
float4 r2 = s_red[2][eg][dm];
float4 r3 = s_red[3][eg][dm];
float4 ba4 = __ldg(&reinterpret_cast<const float4*>(b_a)[0]);
float l0 = (r0.x + r1.x) + (r2.x + r3.x) + ba4.x;
float l1 = (r0.y + r1.y) + (r2.y + r3.y) + ba4.y;
float l2 = (r0.z + r1.z) + (r2.z + r3.z) + ba4.z;
float l3 = (r0.w + r1.w) + (r2.w + r3.w) + ba4.w;
if (t == horizon - 1) {
reinterpret_cast<float4*>(last_logits_out)[e_global_tid] = make_float4(l0, l1, l2, l3);
}
int act = 0;
float best = l0;
if (l1 > best) { best = l1; act = 1; }
if (l2 > best) { best = l2; act = 2; }
if (l3 > best) { best = l3; act = 3; }
if (act == 0) ay = max(0, ay - 1);
else if (act == 1) ay = min(10, ay + 1);
else if (act == 2) ax = max(0, ax - 1);
else if (act == 3) ax = min(10, ax + 1);
bool hit = (ax == fx) && (ay == fy);
if (hit) reward_sum += 1.0f;
rng = lcg_step_dev(rng);
int nfx = (int)(rng % 11ULL);
rng = lcg_step_dev(rng);
int nfy = (int)(rng % 11ULL);
if (hit) {
fx = nfx;
fy = nfy;
}
if (t + 1 < horizon) {
s_obs[tid] = make_float4(
(float)(fx - ax) / 11.0f,
(float)(fy - ay) / 11.0f,
(float)ax / 10.0f,
(float)ay / 10.0f
);
}
}
if (t + 1 < horizon) {
__syncthreads();
}
}
if (tid < M_TILE && e_global_tid < num_envs) {
rewards_out[e_global_tid] = reward_sum;
positions_out[e_global_tid * 2 + 0] = (int64_t)ax;
positions_out[e_global_tid * 2 + 1] = (int64_t)ay;
}
}
extern "C" {
void c_launch_mt19937_init(int64_t seed, int num_envs, int8_t* out_agent_food, cudaStream_t stream) {
mt19937_init_kernel<<<1, 256, 0, stream>>>((uint64_t)seed, num_envs, out_agent_food);
}
void c_launch_env_step(
int N,
const float* agent_in,
const float* food_in,
const int64_t* actions,
const int64_t* rng_in,
float* agent_out,
float* food_out,
float* reward_out,
int64_t* rng_out,
int* any_hit_flag,
cudaStream_t stream
) {
cudaMemsetAsync(any_hit_flag, 0, sizeof(int), stream);
int threads = 256;
int blocks = (N + threads - 1) / threads;
env_step_phase1_kernel<<<blocks, threads, 0, stream>>>(
N, agent_in, food_in, actions, agent_out, reward_out, any_hit_flag
);
env_step_phase2_kernel<<<blocks, threads, 0, stream>>>(
N, food_in, reward_out, rng_in, food_out, rng_out, any_hit_flag
);
}
void c_launch_policy_forward(
int N,
const float* obs,
const float* state_in,
const float* w_enc,
const float* b_enc,
const float* w_gru_T,
const float* w_a,
const float* b_a,
const float* w_v,
const float* b_v,
float* logits_out,
float* state_out,
float* value_out,
cudaStream_t stream
) {
constexpr int G = 4;
constexpr int M_REG = 4;
constexpr int M_TILE = G * M_REG;
dim3 block(64, G);
dim3 grid((N + M_TILE - 1) / M_TILE);
policy_forward_kernel<G, M_REG><<<grid, block, 0, stream>>>(
N, obs, state_in, w_enc, b_enc, w_gru_T, w_a, b_a, w_v, b_v,
logits_out, state_out, value_out
);
}
void c_launch_fused_rollout(
int num_envs,
int horizon,
int64_t seed,
int8_t* init_agent_food,
const float* w_enc,
const float* b_enc,
const float* w_gru_packed,
const float* w_a,
const float* b_a,
float* l2_state,
float* rewards_out,
int64_t* positions_out,
float* last_logits_out,
cudaStream_t stream
) {
mt19937_init_kernel<<<1, 256, 0, stream>>>((uint64_t)seed, num_envs, init_agent_food);
const float4* wp = reinterpret_cast<const float4*>(w_gru_packed);
if (num_envs <= 4096) {
constexpr int G = 2;
constexpr int M_REG = 11;
constexpr int UNROLL_KP = 32;
constexpr int M_TILE = G * M_REG;
constexpr int V_DIM = (M_REG + 3) / 4;
constexpr int ROWS = 4 * M_REG;
constexpr int smem = (G * V_DIM * 256) * sizeof(float4) + ROWS * 256 * sizeof(float);
static bool attr_set_m11 = false;
if (!attr_set_m11) {
cudaFuncSetAttribute(
kernel_c2_g2_any<M_REG, UNROLL_KP>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem
);
attr_set_m11 = true;
}
dim3 block(128, G);
dim3 grid((num_envs + M_TILE - 1) / M_TILE);
kernel_c2_g2_any<M_REG, UNROLL_KP><<<grid, block, smem, stream>>>(
num_envs, horizon, seed, init_agent_food,
w_enc, b_enc, wp, w_a, b_a,
rewards_out, positions_out, last_logits_out
);
} else {
constexpr int G = 2;
constexpr int M_REG = 22;
constexpr int SMEM_ENVS = 24;
constexpr int UNROLL_KP = 16;
constexpr int M_TILE = G * M_REG;
constexpr int NUM_THREADS = G * 128;
constexpr int V_DIM = (M_REG + 3) / 4;
constexpr int ROWS = 2 * SMEM_ENVS;
constexpr int smem = (G * V_DIM * 256) * sizeof(float4) + ROWS * NUM_THREADS * sizeof(float);
static bool attr_set_m22 = false;
if (!attr_set_m22) {
cudaFuncSetAttribute(
kernel_c2_smem_l2_v2<G, M_REG, SMEM_ENVS, UNROLL_KP>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem
);
attr_set_m22 = true;
}
dim3 block(128, G);
dim3 grid((num_envs + M_TILE - 1) / M_TILE);
kernel_c2_smem_l2_v2<G, M_REG, SMEM_ENVS, UNROLL_KP><<<grid, block, smem, stream>>>(
num_envs, horizon, seed, init_agent_food,
w_enc, b_enc, wp, w_a, b_a, l2_state,
rewards_out, positions_out, last_logits_out
);
}
}
} // extern "C"
"""
_CPP_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cstdint>
#include <vector>
extern "C" {
void c_launch_mt19937_init(int64_t seed, int num_envs, int8_t* out_agent_food, cudaStream_t stream);
void c_launch_env_step(
int N,
const float* agent_in,
const float* food_in,
const int64_t* actions,
const int64_t* rng_in,
float* agent_out,
float* food_out,
float* reward_out,
int64_t* rng_out,
int* any_hit_flag,
cudaStream_t stream
);
void c_launch_policy_forward(
int N,
const float* obs,
const float* state_in,
const float* w_enc,
const float* b_enc,
const float* w_gru_T,
const float* w_a,
const float* b_a,
const float* w_v,
const float* b_v,
float* logits_out,
float* state_out,
float* value_out,
cudaStream_t stream
);
void c_launch_fused_rollout(
int num_envs,
int horizon,
int64_t seed,
int8_t* init_agent_food,
const float* w_enc,
const float* b_enc,
const float* w_gru_packed,
const float* w_a,
const float* b_a,
float* l2_state,
float* rewards_out,
int64_t* positions_out,
float* last_logits_out,
cudaStream_t stream
);
}
void launch_mt19937_init(int64_t seed, int num_envs, torch::Tensor out_agent_food) {
c_launch_mt19937_init(seed, num_envs, out_agent_food.data_ptr<int8_t>(), at::cuda::getCurrentCUDAStream());
}
std::vector<torch::Tensor> launch_env_step(
torch::Tensor agent,
torch::Tensor food,
torch::Tensor actions,
torch::Tensor rng_state,
torch::Tensor any_hit_flag
) {
int N = agent.size(0);
auto agent_out = torch::empty_like(agent);
auto food_out = torch::empty_like(food);
auto reward_out = torch::empty({N}, agent.options());
auto rng_out = torch::empty_like(rng_state);
c_launch_env_step(
N,
agent.data_ptr<float>(),
food.data_ptr<float>(),
actions.data_ptr<int64_t>(),
rng_state.data_ptr<int64_t>(),
agent_out.data_ptr<float>(),
food_out.data_ptr<float>(),
reward_out.data_ptr<float>(),
rng_out.data_ptr<int64_t>(),
any_hit_flag.data_ptr<int>(),
at::cuda::getCurrentCUDAStream()
);
return {agent_out, food_out, reward_out, rng_out};
}
std::vector<torch::Tensor> launch_policy_forward(
torch::Tensor obs,
torch::Tensor state,
torch::Tensor w_enc,
torch::Tensor b_enc,
torch::Tensor w_gru_T,
torch::Tensor w_a,
torch::Tensor b_a,
torch::Tensor w_v,
torch::Tensor b_v
) {
int N = obs.size(0);
auto logits_out = torch::empty({N, 4}, obs.options());
auto state_out = torch::empty_like(state);
auto value_out = torch::empty({N}, obs.options());
c_launch_policy_forward(
N,
obs.data_ptr<float>(),
state.data_ptr<float>(),
w_enc.data_ptr<float>(),
b_enc.data_ptr<float>(),
w_gru_T.data_ptr<float>(),
w_a.data_ptr<float>(),
b_a.data_ptr<float>(),
w_v.data_ptr<float>(),
b_v.data_ptr<float>(),
logits_out.data_ptr<float>(),
state_out.data_ptr<float>(),
value_out.data_ptr<float>(),
at::cuda::getCurrentCUDAStream()
);
return {logits_out, state_out, value_out};
}
void launch_fused_rollout(
int num_envs,
int horizon,
int64_t seed,
torch::Tensor init_agent_food,
torch::Tensor w_enc,
torch::Tensor b_enc,
torch::Tensor w_gru_packed,
torch::Tensor w_a,
torch::Tensor b_a,
torch::Tensor l2_state,
torch::Tensor rewards_out,
torch::Tensor positions_out,
torch::Tensor last_logits_out
) {
c_launch_fused_rollout(
num_envs,
horizon,
seed,
init_agent_food.data_ptr<int8_t>(),
w_enc.data_ptr<float>(),
b_enc.data_ptr<float>(),
w_gru_packed.data_ptr<float>(),
w_a.data_ptr<float>(),
b_a.data_ptr<float>(),
l2_state.data_ptr<float>(),
rewards_out.data_ptr<float>(),
positions_out.data_ptr<int64_t>(),
last_logits_out.data_ptr<float>(),
at::cuda::getCurrentCUDAStream()
);
}
"""
_ext = load_inline(
name="grid_mingru_cuda_ext_v9",
cpp_sources=_CPP_SRC,
cuda_sources=_CUDA_SRC,
functions=[
"launch_mt19937_init",
"launch_env_step",
"launch_policy_forward",
"launch_fused_rollout",
],
extra_cuda_cflags=["-O3", "--use_fast_math"],
verbose=False,
)
def _pack_w_gru(w: torch.Tensor) -> torch.Tensor:
wt = w.detach().view(3, 3, 128, 2, 128, 2)
wp = wt.permute(0, 4, 5, 1, 2, 3).contiguous()
hg0 = wp[:, :, 0, 0:2, :, :].permute(0, 1, 3, 2, 4).reshape(3, 128, 128, 4)
p01 = wp[:, :, :, 2, :, :].permute(0, 1, 3, 2, 4).reshape(3, 128, 128, 4)
hg1 = wp[:, :, 1, 0:2, :, :].permute(0, 1, 3, 2, 4).reshape(3, 128, 128, 4)
return torch.stack([hg0, p01, hg1], dim=2).contiguous()
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._cached_w_gru_T: torch.Tensor | None = None
self._cached_w_gru_packed: torch.Tensor | None = None
self._cached_w_gru_ver: int = -1
self._cached_w_gru_ptr: int = 0
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)
self._cached_w_gru_T = None
self._cached_w_gru_packed = None
def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
self._cached_w_gru_T = None
self._cached_w_gru_packed = None
return super().load_state_dict(state_dict, strict=strict, assign=assign)
def _ensure_gru_cache(self) -> None:
w = self.w_gru
if (
self._cached_w_gru_T is None
or self._cached_w_gru_packed is None
or self._cached_w_gru_ver != w._version
or self._cached_w_gru_ptr != w.data_ptr()
or self._cached_w_gru_T.device != w.device
):
wd = w.detach()
self._cached_w_gru_T = wd.view(3, 3, 256, 256).permute(0, 3, 1, 2).contiguous()
self._cached_w_gru_packed = _pack_w_gru(wd)
self._cached_w_gru_ver = w._version
self._cached_w_gru_ptr = w.data_ptr()
def get_w_gru_T(self) -> torch.Tensor:
self._ensure_gru_cache()
return self._cached_w_gru_T
def get_w_gru_packed(self) -> torch.Tensor:
self._ensure_gru_cache()
return self._cached_w_gru_packed
def forward(self, obs: torch.Tensor, state: torch.Tensor):
return policy_forward(self, obs, state)
_any_hit_flag_buf: torch.Tensor | None = None
_init_af_buf: torch.Tensor | None = None
_l2_state_buf: torch.Tensor | None = None
_default_model: Model | None = None
def _get_any_hit_flag(device: torch.device) -> torch.Tensor:
global _any_hit_flag_buf
if _any_hit_flag_buf is None or _any_hit_flag_buf.device != device:
_any_hit_flag_buf = torch.zeros(1, dtype=torch.int32, device=device)
return _any_hit_flag_buf
def _get_init_af_buf(num_envs: int, device: torch.device) -> torch.Tensor:
global _init_af_buf
req = max(num_envs, 65536) * 4
if _init_af_buf is None or _init_af_buf.numel() < req or _init_af_buf.device != device:
_init_af_buf = torch.empty(req, dtype=torch.int8, device=device)
return _init_af_buf
def _get_l2_state_buf(num_envs: int, device: torch.device) -> torch.Tensor:
global _l2_state_buf
max_ctas = (max(num_envs, 65536) + 43) // 44
req = max_ctas * 84 * 256
if _l2_state_buf is None or _l2_state_buf.numel() < req or _l2_state_buf.device != device:
_l2_state_buf = torch.empty(req, dtype=torch.float32, device=device)
return _l2_state_buf
def policy_forward(model: Model, obs: torch.Tensor, state: torch.Tensor):
obs_c = obs.contiguous()
state_c = state.contiguous()
if hasattr(model, "get_w_gru_T"):
w_gru_T = model.get_w_gru_T()
else:
w_gru_T = model.w_gru.detach().view(3, 3, 256, 256).permute(0, 3, 1, 2).contiguous()
logits, new_state, value = _ext.launch_policy_forward(
obs_c,
state_c,
model.w_enc.contiguous(),
model.b_enc.contiguous(),
w_gru_T,
model.w_a.contiguous(),
model.b_a.contiguous(),
model.w_v.contiguous(),
model.b_v.contiguous(),
)
return logits, new_state, value
def env_step(
agent: torch.Tensor,
food: torch.Tensor,
actions: torch.Tensor,
rng_state: torch.Tensor,
):
flag = _get_any_hit_flag(agent.device)
agent_out, food_out, reward_out, rng_out = _ext.launch_env_step(
agent.contiguous(),
food.contiguous(),
actions.contiguous(),
rng_state.contiguous(),
flag,
)
return agent_out, food_out, reward_out, rng_out
def run(num_envs: int, horizon: int, seed: int, model: Model | None = None) -> dict:
global _default_model
device = torch.device("cuda:0")
if model is None:
if _default_model is None:
_default_model = Model().to(device).eval()
model = _default_model
elif next(model.parameters()).device != device:
model = model.to(device)
if num_envs < 512:
af_buf = _get_init_af_buf(num_envs, device)
_ext.launch_mt19937_init(seed, num_envs, af_buf)
agent = af_buf[: num_envs * 2].view(num_envs, 2).float()
food = af_buf[num_envs * 2 : num_envs * 4].view(num_envs, 2).float()
rng_state = torch.arange(num_envs, device=device, dtype=torch.int64) + (seed * 10007)
state = torch.zeros(num_envs, GRU_LAYERS, HIDDEN, device=device, dtype=torch.float32)
rewards = torch.zeros(num_envs, device=device, dtype=torch.float32)
last_logits = torch.zeros(num_envs, NUM_ACTIONS, device=device, dtype=torch.float32)
for _t in range(horizon):
obs = torch.stack(
[
(food[:, 0] - agent[:, 0]) / BOARD,
(food[:, 1] - agent[:, 1]) / BOARD,
agent[:, 0] / (BOARD - 1),
agent[:, 1] / (BOARD - 1),
],
dim=-1,
)
logits, state, _ = policy_forward(model, obs, state)
last_logits = logits
actions = torch.argmax(logits, dim=-1)
agent, food, r, rng_state = env_step(agent, food, actions, rng_state)
rewards = rewards + r
return {
"rewards": rewards,
"positions": agent.round().long(),
"last_logits": last_logits,
}
if hasattr(model, "get_w_gru_packed"):
w_gru_packed = model.get_w_gru_packed()
else:
w_gru_packed = _pack_w_gru(model.w_gru)
af_buf = _get_init_af_buf(num_envs, device)
l2_buf = _get_l2_state_buf(num_envs, device)
rewards = torch.empty(num_envs, dtype=torch.float32, device=device)
positions = torch.empty((num_envs, 2), dtype=torch.int64, device=device)
last_logits = torch.empty((num_envs, NUM_ACTIONS), dtype=torch.float32, device=device)
_ext.launch_fused_rollout(
num_envs,
horizon,
seed,
af_buf,
model.w_enc,
model.b_enc,
w_gru_packed,
model.w_a,
model.b_a,
l2_buf,
rewards,
positions,
last_logits,
)
return {
"rewards": rewards,
"positions": positions,
"last_logits": last_logits,
}
20261013_175345_agy_gemini-4-argon-high_04_grid_mingru_sps