init commit

nanovllm/utils/memory.py

@@ -0,0 +1,14 @@
+import os
+import subprocess
+import torch
+
+
+def get_gpu_memory(device_id: int = 0):
+    torch.cuda.synchronize()
+    result = subprocess.check_output(
+        ['nvidia-smi', '-i', str(device_id), '--query-gpu=memory.total,memory.used,memory.free', '--format=csv,nounits,noheader'],
+        encoding='utf-8'
+    )
+    total_memory, used_memory, free_memory = [int(x) for x in result.strip().split(', ')]
+    return total_memory, used_memory, free_memory
+