[Fix: RL] Snapshot async state-capture outputs before overlap (#34319)

2026-08-14 16:39:35 -07:00
parent 8b4faa3336
commit 90b3db6dd8
1 changed files with 3 additions and 2 deletions
@@ -176,9 +176,10 @@ class BaseTopkCapturer:
            forward_batch, can_run_graph, cuda_graph_batch
        )
        if no_copy_to_cpu:
+            # Clone before the next overlapping forward reuses these buffers.
            return TopkCaptureOutput(
-                out_cache_loc=forward_batch.out_cache_loc,
-                topk=slice_gpu,
+                out_cache_loc=forward_batch.out_cache_loc.clone(),
+                topk=slice_gpu.clone(),
                host_cache=self.host_cache,
            )
        out_cache_loc_cpu = forward_batch.out_cache_loc.cpu()