[Misc] Make mem utils can be reused by other platforms (#32322)

Signed-off-by: shen-shanshan <467638484@qq.com>
This commit is contained in:
Shanshan Shen
2026-01-14 19:46:01 +08:00
committed by GitHub
parent 3f28174c6a
commit ce0946249d
3 changed files with 21 additions and 25 deletions

View File

@@ -312,9 +312,6 @@ class Worker(WorkerBase):
logger.info(msg)
return kv_cache_memory_bytes
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
# Execute a forward pass with dummy inputs to profile the memory usage
# of the model.
with memory_profiling(
@@ -360,7 +357,6 @@ class Worker(WorkerBase):
format_gib(self.available_kv_cache_memory_bytes),
scope="local",
)
gc.collect()
return int(self.available_kv_cache_memory_bytes)