From 0bade503bfb0a517dc6ef4c0ece199b694f30609 Mon Sep 17 00:00:00 2001
From: Valtteri Valo <valtterimiikovalo@gmail.com>
Date: Fri, 29 May 2026 20:37:36 +0300
Subject: [PATCH] Reset recurrent state corrupted by cudagraph warmup

Warmup advances per-buffer and frozen-bank MinGRU state but never restored it, unlike weights, momentum, and RNG state. A fresh PuffeRL now starts from zero recurrent state.
---
 src/pufferlib.cu | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/pufferlib.cu b/src/pufferlib.cu
index df6faa57f4..a6c8c5957a 100644
--- a/src/pufferlib.cu
+++ b/src/pufferlib.cu
@@ -2091,6 +2091,11 @@ std::unique_ptr<PuffeRL> create_pufferl_impl(HypersT& hypers,
         }
         cudaMemset(pufferl->rng_offset_puf.data, 0,
             numel(pufferl->rng_offset_puf.shape) * sizeof(long));
+        for (int bank = 0; bank < 1 + pufferl->num_frozen_banks; bank++) {
+            PrecisionTensor* bs = (bank == 0) ? pufferl->buffer_states
+                                              : pufferl->frozen_banks[bank - 1].buffer_states;
+            for (int b = 0; b < num_buffers; b++) puf_zero(&bs[b], pufferl->default_stream);
+        }
         cudaDeviceSynchronize();
 
         pufferl->epoch = 0;