Unverified Commit 00c4cb56 authored by Matthew Bonanni's avatar Matthew Bonanni Committed by GitHub
Browse files

[Bugfix] Clear stale CG keys after memory profiling (#36416)


Signed-off-by: default avatarMatthew Bonanni <mbonanni@redhat.com>
parent 941e52c2
...@@ -5644,6 +5644,9 @@ class GPUModelRunner( ...@@ -5644,6 +5644,9 @@ class GPUModelRunner(
for instance in list(CUDAGraphWrapper._all_instances): for instance in list(CUDAGraphWrapper._all_instances):
if id(instance) in original_pools: if id(instance) in original_pools:
instance.graph_pool = original_pools[id(instance)] instance.graph_pool = original_pools[id(instance)]
for key_set in self.cudagraph_dispatcher.cudagraph_keys.values():
key_set.clear()
self.cudagraph_dispatcher.keys_initialized = False
self.maybe_remove_all_loras(self.lora_config) self.maybe_remove_all_loras(self.lora_config)
self._cleanup_profiling_kv_cache() self._cleanup_profiling_kv_cache()
compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment