pass static activ buffers

Author: jsuarez5341

pufferlib/extensions/bindings.cpp

@@ -226,7 +226,7 @@ TORCH_LIBRARY_IMPL(pufferlib, CUDA, m) {
 }
 
 TORCH_LIBRARY(_C, m) {
-    m.def("mingru_gate(Tensor state, Tensor combined) -> (Tensor, Tensor)");
+    m.def("mingru_gate(Tensor state, Tensor combined, Tensor out, Tensor next_state) -> ()");
     m.def("fc_max(Tensor x, Tensor W, Tensor b) -> Tensor");
 }
 
@@ -351,8 +351,8 @@ PYBIND11_MODULE(_C, m) {
                          int num_layers, int num_atns, bool continuous) {
             return new Policy(alloc, input, hidden, output, num_layers, num_atns, continuous);
         }))
-        .def("forward", &Policy::forward)
-        .def("forward_train", &Policy::forward_train)
+        .def("forward", static_cast<std::tuple<Logits, Tensor, Tensor> (Policy::*)(Tensor, Tensor)>(&Policy::forward))
+        .def("forward_train", static_cast<std::tuple<Logits, Tensor> (Policy::*)(Tensor, Tensor)>(&Policy::forward_train))
         .def("init_weights", &Policy::init_weights)
         .def("parameters", &Policy::parameters)
         .def("named_parameters", [](Policy& self) {

pufferlib/extensions/cuda/kernels.cu

@@ -638,8 +638,12 @@ __global__ void ppo_loss_backward_kernel_optimized(
     int t = idx % T_seq;
     int nt = n * T_seq + t;
 
+    // Input strides (for reading non-contiguous logits/values_pred)
     int logits_base = n * logits_stride_n + t * logits_stride_t;
     int values_idx = n * values_stride_n + t * values_stride_t;
+    // Output indices (for writing to contiguous grad buffers)
+    int grad_logits_base = nt * A_total;
+    int grad_values_idx = nt;
 
     float old_logp = to_float(old_logprobs[nt]);
     float adv = float(advantages[nt]);
@@ -672,7 +676,7 @@ __global__ void ppo_loss_backward_kernel_optimized(
     } else {
         d_val_pred = val_pred - ret;
     }
-    grad_values_pred[values_idx] = dL * vf_coef * d_val_pred;
+    grad_values_pred[grad_values_idx] = dL * vf_coef * d_val_pred;
 
     if (is_continuous) {
         // Continuous: compute total log prob first for ratio
@@ -724,14 +728,14 @@ __global__ void ppo_loss_backward_kernel_optimized(
 
             // Gradient wrt mean: d_log_prob/d_mean = (action - mean) / var
             float d_mean = d_new_logp * diff / var;
-            grad_logits[logits_base + h * logits_stride_a] = d_mean;
+            grad_logits[grad_logits_base + h] = d_mean;
 
             // Gradient wrt log_std:
             // d_log_prob/d_log_std = (action - mean)^2 / var - 1
             // d_entropy/d_log_std = 1
             // Total: d_new_logp * ((diff^2/var) - 1) + d_entropy_term * 1
             float d_log_std = d_new_logp * (diff * diff / var - 1.0f) + d_entropy_term;
-            grad_logstd[logits_base + h * logits_stride_a] = d_log_std;
+            grad_logstd[grad_logits_base + h] = d_log_std;
         }
     } else {
         // Discrete: original implementation
@@ -822,7 +826,7 @@ __global__ void ppo_loss_backward_kernel_optimized(
                 // Each head's entropy contributes independently to total entropy
                 d_logit += d_entropy_term * p * (-ent - logp);
 
-                grad_logits[logits_base + (logits_offset + a) * logits_stride_a] = d_logit;
+                grad_logits[grad_logits_base + logits_offset + a] = d_logit;
             }
 
             logits_offset += A;

pufferlib/extensions/cuda/modules.cu

@@ -25,7 +25,7 @@ using AutogradCtx = torch::autograd::AutogradContext;
 // returns {out, next_state} where:
 //   out (B, H) = sigmoid(proj) * mingru_out
 //   next_state (B, H) = mingru_out (for recurrence)
-vector<Tensor> mingru_gate(Tensor state, Tensor combined) {
+void mingru_gate(Tensor state, Tensor combined, Tensor out, Tensor next_state) {
     TORCH_CHECK(state.is_cuda(), "state must be on CUDA");
     TORCH_CHECK(combined.is_cuda(), "combined must be on CUDA");
     TORCH_CHECK(state.dtype() == combined.dtype(), "dtypes must match");
@@ -36,9 +36,6 @@ vector<Tensor> mingru_gate(Tensor state, Tensor combined) {
 
     int B = static_cast<int>(state.size(0));
     int H = static_cast<int>(state.size(1));
-
-    auto out = torch::empty_like(state);
-    auto next_state = torch::empty_like(state);
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
     mingru_gate_inference_kernel<<<grid_size(B * H), BLOCK_SIZE, 0, stream>>>(
@@ -47,7 +44,6 @@ vector<Tensor> mingru_gate(Tensor state, Tensor combined) {
         (const precision_t*)combined.data_ptr(),
         (const precision_t*)state.data_ptr(),
         H, B);
-    return {out, next_state};
 }
 
 // PrefixScan: checkpointed associative scan for MinGRU training
@@ -452,16 +448,16 @@ void ppo_loss_fwd_bwd(
     TORCH_CHECK(act_sizes.is_cuda() && act_sizes.dtype() == torch::kInt32,
                 "act_sizes must be int32 on CUDA");
 
-    // Make inputs contiguous for both kernels
-    logits = logits.contiguous();
-    values_pred = values_pred.contiguous();
-    old_logprobs = old_logprobs.contiguous();
-    advantages = advantages.contiguous();
-    prio = prio.contiguous();
-    values = values.contiguous();
-    returns = returns.contiguous();
+    // logits/values_pred may be non-contiguous (fused decoder output) — kernel handles via strides
+    // Grad outputs use contiguous layout (nt * A_total indexing)
+    TORCH_CHECK(old_logprobs.is_contiguous(), "old_logprobs must be contiguous");
+    TORCH_CHECK(advantages.is_contiguous(), "advantages must be contiguous");
+    TORCH_CHECK(prio.is_contiguous(), "prio must be contiguous");
+    TORCH_CHECK(values.is_contiguous(), "values must be contiguous");
+    TORCH_CHECK(returns.is_contiguous(), "returns must be contiguous");
 
     bool is_continuous = logstd.defined() && logstd.numel() > 0;
+    // TODO: pre-allocate contiguous logstd buffer to remove this alloc
     if (is_continuous) logstd = logstd.contiguous();
 
     int N = static_cast<int>(logits.size(0));
@@ -471,7 +467,8 @@ void ppo_loss_fwd_bwd(
     int total = N * T;
 
     auto [adv_var, adv_mean] = torch::var_mean(advantages);
-    auto actions_flat = actions.reshape({total, num_atns}).contiguous();
+    auto actions_flat = actions.reshape({total, num_atns});
+    TORCH_CHECK(actions_flat.is_contiguous(), "actions must be contiguous");
 
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();