KernelBench cuda · RTX PRO 6000
Grid + MinGRU SPS DeepSeek V4 Flash (0731)
14.5%geomean peak fraction across shapes
manually audited: clean
Manual static audit covered all 133 lines of solution.py, all 450 lines of scratch/kernels.cu, all 243287 transcript events and 272 agent tool calls, result.json, check.log, benchmark.log, the CUDA-language sidecar, and the frozen grader/template files. solution.py:85-98 takes the current model parameters and current obs/state, while solution.py:111-127 generates the requested seed's initial state and calls the raw-CUDA rollout. The CUDA sidecar performs real input- and weight-dependent work: kernels.cu:59-225 computes the encoder, every gate of all three MinGRU layers, logits, value, and greedy actions; kernels.cu:228-271 updates positions, rewards, food, and LCG state; kernels.cu:285-294 permutes the supplied live GRU weights; and kernels.cu:387-443 executes every requested horizon step and returns newly allocated rewards, positions, and logits. No forbidden library, Triton/DSL, reference import, stack/caller or check.py sniffing, result-file read, constant/seed answer table, output memoization, CUDA graph, data_ptr identity cache, or persistent computed output exists. The only static cache is the one-time CUDA function shared-memory attribute in kernels.cu:299-309; it cannot cache an answer. Therefore no empirical same-buffer overwrite cache test is required. The transcript reads only this cell's grader and eval sources, writes only this cell's solution/kernels plus disposable /tmp experiments, and contains no foreign-run artifact access. All seven archived repo grader/template files compare byte-for-byte equal to template_files, consistent with result.json template_mutated=false. The official check.log reports cuda_language ok, framework=cuda_raw, then PASS; check.py runs its numeric_stress_cases loop, but check.log gives no per-case magnitudes, so none are claimed here. cuda_language.json reports triton_cheat=false, dsl_cheat=false, forbidden_hits=[], and ok=true. The publishable result is the sequential isolated regrade recorded in result.json at 2026-08-03T06:31:29+00:00 on host brev-6h4h8jy86, GPU index 0, an NVIDIA RTX PRO 6000 Blackwell Server Edition. The isolated check passed with exit code zero, and benchmark.log records solution SPS 14838607.270, 24783601.782, 30218454.802, and 20302314.303 for shapes 0-3, peak_fraction=0.1453, RESULT: OK, and benchmark exit code zero. result.json records correct=true and regrade mode sequential_isolated. Its regrade.contended record preserves the pre-regrade evidence: correct=true, peak_fraction=0.1419, and zero check/benchmark exit codes. The annotation publishes only the isolated 0.1453 metric; 0.1419 remains provenance, not the publication score.
harnessor-fableagent session1h 47mtotal wall1h 48mcheck33sbenchmark2soutput tokens—gpu-lock wait50sgpu-lock held36sregimethroughput
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)
"""CUDA grid-foraging + 3x MinGRU(h=256) policy. Fast rollout on RTX PRO 6000.
Exposes Model / policy_forward / env_step / run matching reference.py exactly.
The heavy lifting is a real CUDA C++ kernel (kernels.cu) compiled via
torch.utils.cpp_extension.load.
"""
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
_HERE = Path(__file__).resolve().parent
def _setup_env() -> None:
# Prefer a real nvcc over any shim; ensure ninja is findable.
cur = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
if not (cur and (Path(cur) / "bin/nvcc").exists()):
for cand in ("/usr/local/cuda-12.8", "/usr/local/cuda", "/usr/local/cuda-13.0"):
if (Path(cand) / "bin/nvcc").exists():
os.environ["CUDA_HOME"] = cand
break
# Locate ninja (torch's cpp_extension requires it).
cand_bins = []
if os.environ.get("VIRTUAL_ENV"):
cand_bins.append(str(Path(os.environ["VIRTUAL_ENV"]) / "bin"))
try:
import ninja # type: ignore
cand_bins.append(ninja.BIN_DIR)
except Exception:
pass
for cand_bin in cand_bins:
if os.path.isdir(cand_bin) and os.path.exists(os.path.join(cand_bin, "ninja")):
os.environ["PATH"] = cand_bin + os.pathsep + os.environ.get("PATH", "")
break
_setup_env()
from torch.utils.cpp_extension import load as _ext_load # noqa: E402
_ext = _ext_load(
name="grid_mingru_ext",
sources=[str(_HERE / "kernels.cu")],
extra_cuda_cflags=["-O3", "--use_fast_math"],
verbose=False,
)
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)
def _weights(model: nn.Module, device=None):
dev = device or next(model.parameters()).device
return tuple(
p.detach().to(device=dev, dtype=torch.float32).contiguous()
for p in (model.w_enc, model.b_enc, model.w_gru, model.w_a,
model.b_a, model.w_v, model.b_v)
)
def policy_forward(model: nn.Module, obs: torch.Tensor, state: torch.Tensor):
w = _weights(model, obs.device)
logits, state_out, value = _ext.policy_forward(*w, obs.contiguous(),
state.contiguous())
return logits, state_out, value
def env_step(agent: torch.Tensor, food: torch.Tensor,
actions: torch.Tensor, rng_state: torch.Tensor):
agent = agent.contiguous()
food = food.contiguous()
actions = actions.contiguous()
rng_state = rng_state.contiguous()
agent_o, food_o, reward, rng_o = _ext.env_step(agent, food, actions, rng_state)
return agent_o, food_o, reward, rng_o
def run(num_envs: int, horizon: int, seed: int, model: Model | None = None) -> dict:
device = torch.device("cuda:0")
if model is None:
model = Model()
model = model.to(device).eval()
w = _weights(model, device)
# Initial conditions must match reference.run exactly (CPU MT19937).
g = torch.Generator(device="cpu")
g.manual_seed(seed)
agent = torch.randint(0, BOARD, (num_envs, 2), generator=g)
food = torch.randint(0, BOARD, (num_envs, 2), generator=g)
rng_state = torch.arange(num_envs, device=device, dtype=torch.int64) + (seed * 10007)
rewards, positions, last_logits = _ext.rollout(
*w, agent, food, rng_state, int(num_envs), int(horizon)
)
return {
"rewards": rewards,
"positions": positions,
"last_logits": last_logits,
}
# ==================================================================
# ===== sidecar: kernels.cu (18907 bytes, loaded by solution.py) =====
# ==================================================================
// Grid-foraging + 3x MinGRU(h=256) CUDA kernels.
// Policy is a block-tiled SGEMM: each block handles BN envs, keeping the
// recurrent activation h in shared memory between layers.
#include <cuda_runtime.h>
#include <cuda_pipeline.h>
#include <cstdint>
#include <torch/extension.h>
#define BOARD 11
#define HIDDEN 256
#define GRU_LAYERS 3
#define GRU_OUT 768
#define NUM_ACTIONS 4
#define LCG_A 6364136223846793005ULL
#define LCG_MASK 0x7FFFFFFFFFFFFFFFULL
// ---- tiling config ----
#define BN 32 // envs per block
#define BM 192 // gate rows per M-pass
#define THREADS 256
#define TM 12 // M-rows per thread (multiple of 3 -> TM/3 hidden units)
#define TN 2 // envs per thread
#define KC 8 // K-chunk
#define N_MPASS (GRU_OUT / BM) // 4
#define N_KC (HIDDEN / KC) // 32
#define N_HU (TM / 3) // hidden units per thread per pass
#define SHMEM_WS_OFF (HIDDEN * BN) // floats (single h buffer)
#define SHMEM_OBS_OFF (SHMEM_WS_OFF + 2 * KC * BM) // floats (double-buffered Ws)
#define SHMEM_TOTAL_FLOATS (SHMEM_OBS_OFF + BN * 4)
__device__ __forceinline__ float sigmoidf(float x) { return 1.0f / (1.0f + expf(-x)); }
__device__ __forceinline__ uint64_t lcg_step(uint64_t r) {
return (r * LCG_A + 1ULL) & LCG_MASK;
}
__device__ __forceinline__ void stage_ws(float* Ws, int buf, const float* wlT,
int m_base, int k0) {
// copy KC*BM floats (16B chunks) from global wlT into shared Ws[buf]
constexpr int n4 = (KC * BM) >> 2;
for (int i = threadIdx.x; i < n4; i += THREADS) {
int kk = i / (BM >> 2);
int mm4 = i % (BM >> 2);
const float* src = wlT + (k0 + kk) * GRU_OUT + m_base + (mm4 * 4);
float4* dst = reinterpret_cast<float4*>(Ws + buf * (KC * BM) + kk * BM + (mm4 * 4));
__pipeline_memcpy_async(dst, src, 16);
}
__pipeline_commit();
}
// Block-tiled MLP. smem layout:
// [0, HIDDEN*BN) : h buffer (k-major: [k][n])
// [SHMEM_WS_OFF, +2*KC*BM): double-buffered W staging chunk (k-major: [kk][m])
// [SHMEM_OBS_OFF, +BN*4) : obs buffer [n][i]
// mode: 0 = obs direct, 1 = obs from agent/food
__global__ void __launch_bounds__(THREADS) policy_tiled_kernel(
const float* __restrict__ w_enc, const float* __restrict__ b_enc,
const float* __restrict__ w_permT, // (3, 256, 768) permuted-transposed GRU weights
const float* __restrict__ w_a, const float* __restrict__ b_a,
const float* __restrict__ w_v, const float* __restrict__ b_v,
const float* __restrict__ obs_in, // (N,4) for mode 0
const int* __restrict__ agent, const int* __restrict__ food, // for mode 1
const float* __restrict__ state_in, float* __restrict__ state_out,
float* __restrict__ logits_out, float* __restrict__ value_out,
int* __restrict__ actions_out, int N, int mode)
{
extern __shared__ float smem[];
float* buf0 = smem; // [HIDDEN*BN] (h, k-major [k][n])
float* Ws = smem + SHMEM_WS_OFF; // [2*KC*BM]
float* obs_s = smem + SHMEM_OBS_OFF; // [BN*4]
const int tid = threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5; // 0..7
const int warp_m = warp >> 1; // 0..3
const int warp_n = warp & 1; // 0..1
const int thread_m = lane >> 3; // 0..3
const int thread_n = lane & 7; // 0..7
const int env_base = blockIdx.x * BN;
// m offset (relative to M-pass base) and env offset (relative to block envs)
const int m_off = warp_m * (BM / 4) + thread_m * TM; // 0..BM
const int env_off = warp_n * (BN / 2) + thread_n * TN; // 0..BN
// ---- load obs into shared ----
#pragma unroll
for (int idx = tid; idx < BN * 4; idx += THREADS) {
int n = idx >> 2, i = idx & 3;
int n_abs = env_base + n;
if (n_abs < N) {
if (mode == 0) {
obs_s[n * 4 + i] = obs_in[n_abs * 4 + i];
} else {
int ax = agent[n_abs * 2], ay = agent[n_abs * 2 + 1];
int fx = food[n_abs * 2], fy = food[n_abs * 2 + 1];
if (i == 0) obs_s[n * 4 + 0] = (float)(fx - ax) / (float)BOARD;
else if (i == 1) obs_s[n * 4 + 1] = (float)(fy - ay) / (float)BOARD;
else if (i == 2) obs_s[n * 4 + 2] = (float)ax / 10.0f;
else obs_s[n * 4 + 3] = (float)ay / 10.0f;
}
}
}
__syncthreads();
// ---- enc: hT[u][n] = b_enc[u] + sum_i obs_s[n*4+i]*w_enc[u*4+i] ----
#pragma unroll
for (int idx = tid; idx < HIDDEN * BN; idx += THREADS) {
int u = idx / BN, n = idx % BN;
int n_abs = env_base + n;
float s = b_enc[u];
if (n_abs < N) {
const float* o = obs_s + n * 4;
const float* we = w_enc + u * 4;
s += we[0] * o[0] + we[1] * o[1] + we[2] * o[2] + we[3] * o[3];
}
buf0[u * BN + n] = s;
}
__syncthreads();
// ---- GRU layers ----
#pragma unroll
for (int l = 0; l < GRU_LAYERS; ++l) {
const float* wlT = w_permT + l * HIDDEN * GRU_OUT; // (256, 768)
float acc[TM][TN];
float nh[N_MPASS * N_HU * TN]; // new hidden values, written back after all passes
#pragma unroll
for (int mp = 0; mp < N_MPASS; ++mp) {
const int m_base = mp * BM;
// zero accumulators
#pragma unroll
for (int tm = 0; tm < TM; ++tm)
#pragma unroll
for (int tn = 0; tn < TN; ++tn) acc[tm][tn] = 0.0f;
// pipelined K-loop: double-buffered Ws via cp.async
stage_ws(Ws, 0, wlT, m_base, 0);
for (int kc = 0; kc < N_KC; ++kc) {
const int k0 = kc * KC;
if (kc + 1 < N_KC) stage_ws(Ws, (kc + 1) & 1, wlT, m_base, k0 + KC);
__pipeline_wait_prior(kc + 1 < N_KC ? 1 : 0);
__syncthreads();
const float* Wp = Ws + (kc & 1) * (KC * BM);
#pragma unroll
for (int kk = 0; kk < KC; ++kk) {
float h0 = buf0[(k0 + kk) * BN + env_off + 0];
float h1 = buf0[(k0 + kk) * BN + env_off + 1];
const float* wp = Wp + kk * BM + m_off;
#pragma unroll
for (int tm = 0; tm < TM; ++tm) {
float w = wp[tm];
acc[tm][0] += w * h0;
acc[tm][1] += w * h1;
}
}
}
// ---- GRU update for this M-pass: keep new h in registers ----
const int u_base = (m_base + m_off) / 3;
#pragma unroll
for (int ul = 0; ul < N_HU; ++ul) {
const int u = u_base + ul;
#pragma unroll
for (int tn = 0; tn < TN; ++tn) {
int n_abs = env_base + env_off + tn;
if (n_abs >= N) continue;
float zh = acc[ul * 3 + 0][tn];
float zg = acc[ul * 3 + 1][tn];
float zp = acc[ul * 3 + 2][tn];
float st = state_in[n_abs * GRU_LAYERS * HIDDEN + l * HIDDEN + u];
float out = st + sigmoidf(zg) * (tanhf(zh) - st);
float p = sigmoidf(zp);
float hn = p * out + (1.0f - p) * buf0[u * BN + env_off + tn];
nh[(mp * N_HU + ul) * TN + tn] = hn;
state_out[n_abs * GRU_LAYERS * HIDDEN + l * HIDDEN + u] = out;
}
}
}
// write new hidden back to buf0 (all passes done; buf0 reads complete)
#pragma unroll
for (int mp = 0; mp < N_MPASS; ++mp) {
const int u_base = (mp * BM + m_off) / 3;
#pragma unroll
for (int ul = 0; ul < N_HU; ++ul) {
const int u = u_base + ul;
#pragma unroll
for (int tn = 0; tn < TN; ++tn) {
buf0[u * BN + env_off + tn] = nh[(mp * N_HU + ul) * TN + tn];
}
}
}
__syncthreads(); // new-h writes visible before next layer reads
}
// ---- output: logits, value, argmax ----
#pragma unroll
for (int tn = 0; tn < TN; ++tn) {
int n_abs = env_base + env_off + tn;
if (n_abs >= N) continue;
float lp[4];
#pragma unroll
for (int a = 0; a < NUM_ACTIONS; ++a) {
float s = b_a[a];
const float* wa = w_a + a * HIDDEN;
#pragma unroll 4
for (int k = 0; k < HIDDEN; ++k) s += wa[k] * buf0[k * BN + env_off + tn];
lp[a] = s;
}
float v = b_v[0];
#pragma unroll 4
for (int k = 0; k < HIDDEN; ++k) v += w_v[k] * buf0[k * BN + env_off + tn];
int best = 0;
#pragma unroll
for (int a = 1; a < NUM_ACTIONS; ++a) if (lp[a] > lp[best]) best = a;
logits_out[n_abs * NUM_ACTIONS + 0] = lp[0];
logits_out[n_abs * NUM_ACTIONS + 1] = lp[1];
logits_out[n_abs * NUM_ACTIONS + 2] = lp[2];
logits_out[n_abs * NUM_ACTIONS + 3] = lp[3];
value_out[n_abs] = v;
actions_out[n_abs] = best;
}
}
// ---- env movement + hit detection + reward accumulation ----
__global__ void env_move_kernel(
const int* __restrict__ agent, const int* __restrict__ food,
const int* __restrict__ actions, int* __restrict__ agent_out,
unsigned char* __restrict__ hit, int* __restrict__ hit_any,
float* __restrict__ rewards, int N)
{
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= N) return;
int a = actions[idx];
int ax = agent[idx * 2], ay = agent[idx * 2 + 1];
int nx = ax, ny = ay;
if (a == 0) { if (ny > 0) --ny; }
else if (a == 1) { if (ny < BOARD - 1) ++ny; }
else if (a == 2) { if (nx > 0) --nx; }
else { if (nx < BOARD - 1) ++nx; }
agent_out[idx * 2] = nx;
agent_out[idx * 2 + 1] = ny;
int h = (nx == food[idx * 2] && ny == food[idx * 2 + 1]) ? 1 : 0;
hit[idx] = (unsigned char)h;
if (h) { atomicOr(hit_any, 1); rewards[idx] += 1.0f; }
}
// ---- respawn: multi-block; if any hit, advance all rng and respawn hit foods ----
__global__ void respawn_kernel(
const unsigned char* __restrict__ hit, const int* __restrict__ hit_any,
uint64_t* __restrict__ rng, int* __restrict__ food, int N)
{
int any = *hit_any;
if (!any) return;
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= N) return;
uint64_t r = lcg_step(rng[idx]);
uint64_t r2 = lcg_step(r);
rng[idx] = r2;
if (hit[idx]) {
food[idx * 2] = (int)(r % BOARD);
food[idx * 2 + 1] = (int)(r2 % BOARD);
}
}
// ---- reset the any-hit flag (1 thread) ----
__global__ void reset_flag_kernel(int* __restrict__ hit_any) {
if (threadIdx.x == 0) *hit_any = 0;
}
// ===================== host helpers =====================
static void check_cuda(const char* tag) {
cudaError_t e = cudaGetLastError();
TORCH_CHECK(e == cudaSuccess, tag, ": ", cudaGetErrorString(e));
}
static const float* wptr(torch::Tensor& t) { return t.data_ptr<float>(); }
static const int* iptr(torch::Tensor& t) { return t.data_ptr<int>(); }
static float* wptr_nc(torch::Tensor& t) { return t.data_ptr<float>(); }
static int* iptr_nc(torch::Tensor& t) { return t.data_ptr<int>(); }
// w_gru (3,768,256) -> w_permT (3,256,768): w_permT[l][k][m] = w_gru[l][(m%3)*256 + m/3][k]
static torch::Tensor permute_gru(torch::Tensor w_gru) {
w_gru = w_gru.contiguous();
auto opts = w_gru.options();
auto out = torch::empty({GRU_LAYERS, HIDDEN, GRU_OUT}, opts);
auto idx = torch::arange(GRU_OUT, w_gru.options().dtype(torch::kInt64));
auto orig = (idx % 3) * HIDDEN + (idx / 3);
auto permuted = w_gru.index_select(1, orig.to(opts.dtype(torch::kLong)));
out.copy_(permuted.permute({0, 2, 1}));
return out;
}
static const int SHMEM_BYTES = SHMEM_TOTAL_FLOATS * 4;
static void ensure_smem(int shmem_bytes) {
static bool configured = false;
if (!configured) {
cudaFuncAttributes attr;
cudaFuncGetAttributes(&attr, policy_tiled_kernel); // init attrs
cudaGetLastError();
cudaError_t e = cudaFuncSetAttribute(
policy_tiled_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_BYTES);
TORCH_CHECK(e == cudaSuccess, "cudaFuncSetAttribute: ", cudaGetErrorString(e));
configured = true;
}
}
static void launch_policy(
torch::Tensor& w_enc, torch::Tensor& b_enc, torch::Tensor& w_permT,
torch::Tensor& w_a, torch::Tensor& b_a, torch::Tensor& w_v, torch::Tensor& b_v,
torch::Tensor& obs, torch::Tensor& agent, torch::Tensor& food,
torch::Tensor& state_in, torch::Tensor& state_out,
torch::Tensor& logits, torch::Tensor& value, torch::Tensor& actions,
int mode)
{
int N = (int)state_in.size(0);
ensure_smem(SHMEM_BYTES);
int blocks = (N + BN - 1) / BN;
policy_tiled_kernel<<<blocks, THREADS, SHMEM_BYTES>>>(
wptr(w_enc), wptr(b_enc), wptr(w_permT),
wptr(w_a), wptr(b_a), wptr(w_v), wptr(b_v),
mode == 0 ? wptr(obs) : nullptr,
mode == 1 ? iptr(agent) : nullptr, mode == 1 ? iptr(food) : nullptr,
wptr(state_in), wptr_nc(state_out),
wptr_nc(logits), wptr_nc(value), iptr_nc(actions), N, mode);
check_cuda("policy_tiled");
}
std::vector<torch::Tensor> policy_forward(
torch::Tensor w_enc, torch::Tensor b_enc, torch::Tensor w_gru,
torch::Tensor w_a, torch::Tensor b_a, torch::Tensor w_v, torch::Tensor b_v,
torch::Tensor obs, torch::Tensor state)
{
auto N = obs.size(0);
auto w_permT = permute_gru(w_gru);
auto logits = torch::empty({N, NUM_ACTIONS}, obs.options());
auto state_out = torch::empty({N, GRU_LAYERS, HIDDEN}, obs.options());
auto value = torch::empty({N}, obs.options());
auto actions = torch::empty({N}, obs.options().dtype(torch::kInt32));
auto obs_c = obs.contiguous();
auto state_c = state.contiguous();
auto dummy = torch::empty({0}, obs.options());
launch_policy(w_enc, b_enc, w_permT, w_a, b_a, w_v, b_v,
obs_c, dummy, dummy, state_c, state_out,
logits, value, actions, 0);
return {logits, state_out, value};
}
// env_step on float tensors -> tuple (agent, food, reward, rng)
std::vector<torch::Tensor> env_step(
torch::Tensor agent_f, torch::Tensor food_f, torch::Tensor actions,
torch::Tensor rng)
{
auto N = agent_f.size(0);
auto agent_i = agent_f.to(torch::kInt32).contiguous();
auto food_i = food_f.to(torch::kInt32).contiguous();
auto acts_i = actions.to(torch::kInt32).contiguous();
auto agent_o = torch::empty_like(agent_i);
auto hit = torch::zeros({N}, agent_f.options().dtype(torch::kUInt8));
auto hit_any = torch::zeros({1}, agent_f.options().dtype(torch::kInt32));
auto reward = torch::zeros({N}, agent_f.options());
auto rng_c = rng.to(torch::kInt64).contiguous();
int n = (int)N;
int threads = 256;
int blocks = (n + threads - 1) / threads;
env_move_kernel<<<blocks, threads>>>(
agent_i.data_ptr<int>(), food_i.data_ptr<int>(), acts_i.data_ptr<int>(),
agent_o.data_ptr<int>(), hit.data_ptr<unsigned char>(), hit_any.data_ptr<int>(),
reward.data_ptr<float>(), n);
check_cuda("env_move");
uint64_t* rng_ptr = reinterpret_cast<uint64_t*>(rng_c.data_ptr<int64_t>());
respawn_kernel<<<blocks, threads>>>(hit.data_ptr<unsigned char>(), hit_any.data_ptr<int>(),
rng_ptr, food_i.data_ptr<int>(), n);
check_cuda("respawn");
reset_flag_kernel<<<1, 1>>>(hit_any.data_ptr<int>());
check_cuda("reset");
auto agent_f_out = agent_o.to(torch::kFloat32);
auto food_f_out = food_i.to(torch::kFloat32);
return {agent_f_out, food_f_out, reward, rng_c};
}
// Full rollout. Returns (rewards, positions, last_logits).
std::vector<torch::Tensor> rollout(
torch::Tensor w_enc, torch::Tensor b_enc, torch::Tensor w_gru,
torch::Tensor w_a, torch::Tensor b_a, torch::Tensor w_v, torch::Tensor b_v,
torch::Tensor agent_cpu, torch::Tensor food_cpu, torch::Tensor rng,
int64_t num_envs, int64_t horizon)
{
int N = (int)num_envs;
int T = (int)horizon;
auto opts = torch::TensorOptions().device(torch::kCUDA);
auto w_permT = permute_gru(w_gru);
auto agent = agent_cpu.to(torch::kCUDA).to(torch::kInt32).contiguous();
auto food = food_cpu.to(torch::kCUDA).to(torch::kInt32).contiguous();
auto rng_i = rng.to(torch::kCUDA).contiguous();
uint64_t* rng_ptr = reinterpret_cast<uint64_t*>(rng_i.data_ptr<int64_t>());
auto state = torch::zeros({N, GRU_LAYERS * HIDDEN}, opts.dtype(torch::kFloat32));
auto state2 = torch::zeros({N, GRU_LAYERS * HIDDEN}, opts.dtype(torch::kFloat32));
auto logits = torch::empty({N, NUM_ACTIONS}, opts.dtype(torch::kFloat32));
auto value = torch::empty({N}, opts.dtype(torch::kFloat32));
auto actions = torch::empty({N}, opts.dtype(torch::kInt32));
auto hit = torch::zeros({N}, opts.dtype(torch::kUInt8));
auto hit_any = torch::zeros({1}, opts.dtype(torch::kInt32));
auto rewards = torch::zeros({N}, opts.dtype(torch::kFloat32));
auto agent_tmp = torch::empty_like(agent);
ensure_smem(SHMEM_BYTES);
int blocks = (N + BN - 1) / BN;
for (int t = 0; t < T; ++t) {
policy_tiled_kernel<<<blocks, THREADS, SHMEM_BYTES>>>(
wptr(w_enc), wptr(b_enc), wptr(w_permT),
wptr(w_a), wptr(b_a), wptr(w_v), wptr(b_v),
nullptr, agent.data_ptr<int>(), food.data_ptr<int>(),
state.data_ptr<float>(), state2.data_ptr<float>(),
logits.data_ptr<float>(), value.data_ptr<float>(),
actions.data_ptr<int>(), N, 1);
check_cuda("policy");
env_move_kernel<<<blocks, 256>>>(
agent.data_ptr<int>(), food.data_ptr<int>(), actions.data_ptr<int>(),
agent_tmp.data_ptr<int>(), hit.data_ptr<unsigned char>(),
hit_any.data_ptr<int>(), rewards.data_ptr<float>(), N);
check_cuda("move");
respawn_kernel<<<blocks, 256>>>(
hit.data_ptr<unsigned char>(), hit_any.data_ptr<int>(),
rng_ptr, food.data_ptr<int>(), N);
check_cuda("respawn");
reset_flag_kernel<<<1, 1>>>(hit_any.data_ptr<int>());
check_cuda("reset");
std::swap(state, state2);
std::swap(agent, agent_tmp);
}
auto positions = agent.to(torch::kInt64);
return {rewards, positions, logits};
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("policy_forward", &policy_forward, "policy forward");
m.def("env_step", &env_step, "env step");
m.def("rollout", &rollout, "rollout");
}
20260802_203919_or-fable_deepseek_deepseek-v4-flash-0731_04_grid_mingru_sps