Skip to content

Commit 8f8fede

Browse files
committed
Self-review: correctly clean up CUDA graphs
1 parent a31c262 commit 8f8fede

1 file changed

Lines changed: 3 additions & 0 deletions

File tree

megatron/core/inference/engines/dynamic_engine.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -725,6 +725,9 @@ def suspend(self):
725725
and not self.context.static_kv_memory_pointers
726726
):
727727
delete_cuda_graphs()
728+
if self.controller._sampling_backend == "flashinfer":
729+
self.controller._sampling._sampling_cuda_graphs.clear()
730+
self.controller._sampling._verification_cuda_graphs.clear()
728731

729732
# Build the list of requests to re-add on resume.
730733
# All waiting requests are always included; active requests are included

0 commit comments

Comments
 (0)