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

This commit is contained in:
Jason Mancuso
2026-08-14 16:39:35 -07:00
committed by GitHub
parent 8b4faa3336
commit 90b3db6dd8
+3 -2
View File
@@ -176,9 +176,10 @@ class BaseTopkCapturer:
forward_batch, can_run_graph, cuda_graph_batch forward_batch, can_run_graph, cuda_graph_batch
) )
if no_copy_to_cpu: if no_copy_to_cpu:
# Clone before the next overlapping forward reuses these buffers.
return TopkCaptureOutput( return TopkCaptureOutput(
out_cache_loc=forward_batch.out_cache_loc, out_cache_loc=forward_batch.out_cache_loc.clone(),
topk=slice_gpu, topk=slice_gpu.clone(),
host_cache=self.host_cache, host_cache=self.host_cache,
) )
out_cache_loc_cpu = forward_batch.out_cache_loc.cpu() out_cache_loc_cpu = forward_batch.out_cache_loc.cpu()