[diffusion] logging: log available gpu mem while loading and generating (#15936)

python/sglang/multimodal_gen/runtime/platforms/cpu.py

@@ -4,7 +4,10 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cpu.py
 
 import platform
+from functools import lru_cache
+from typing import Any
 
+import psutil
 import torch
 
 from sglang.multimodal_gen.runtime.platforms.interface import (
@@ -40,10 +43,10 @@ class CpuPlatform(Platform):
         return platform.machine()
 
     @classmethod
+    @lru_cache(maxsize=1)
     def get_device_total_memory(cls, device_id: int = 0) -> int:
-        # This is a rough estimate for CPU memory
-        # In practice, you might want to use psutil or similar
-        return 0
+
+        return psutil.virtual_memory().total
 
     @classmethod
     def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
@@ -56,6 +59,30 @@ class CpuPlatform(Platform):
         # For CPU, we can't easily get memory usage without additional libraries
         return 0.0
 
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int = 0,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+
+        total_free_memory = psutil.virtual_memory().available
+        # For simplicity, we assume 1 NUMA node for now in this platform abstraction
+        # as get_cpu_ids_by_node is not available in multimodal_gen.runtime.utils
+        n_numa_node = 1
+        free_memory = total_free_memory / n_numa_node
+
+        if distributed:
+            import torch.distributed as dist
+
+            tensor = torch.tensor(free_memory, dtype=torch.float32)
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group)
+            free_memory = float(tensor.item())
+
+        return free_memory / (1 << 30)
+
     @classmethod
     def get_device_communicator_cls(cls) -> str:
         return "sglang.multimodal_gen.runtime.distributed.device_communicators.cpu_communicator.CpuCommunicator"

python/sglang/multimodal_gen/runtime/platforms/cuda.py

@@ -5,12 +5,12 @@
 """Code inside this file can safely assume cuda platform, e.g. importing
 pynvml. However, it should not initialize cuda context.
 """
-
 import os
 from collections.abc import Callable
 from functools import lru_cache, wraps
-from typing import TypeVar
+from typing import Any, TypeVar
 
+import psutil
 import torch
 from typing_extensions import ParamSpec
 
@@ -82,6 +82,7 @@ class CudaPlatformBase(Platform):
         raise NotImplementedError
 
     @classmethod
+    @lru_cache(maxsize=1)
     def get_device_total_memory(cls, device_id: int = 0) -> int:
         raise NotImplementedError
 
@@ -111,6 +112,38 @@ class CudaPlatformBase(Platform):
         torch.cuda.reset_peak_memory_stats(device)
         return float(torch.cuda.max_memory_allocated(device))
 
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int = 0,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+        if empty_cache:
+            torch.cuda.empty_cache()
+
+        # Orin, Thor, Spark
+        # SM 8.7 is Orin, 11.0 is Thor, 12.1 is Spark
+        SHARED_SYSMEM_DEVICE_MEM_SMS = (87, 110, 121)
+        capability = cls.get_device_capability(device_id)
+        sm = capability.to_int() if capability else 0
+
+        if sm in SHARED_SYSMEM_DEVICE_MEM_SMS:
+
+            free_gpu_memory = psutil.virtual_memory().available
+        else:
+            free_gpu_memory, _ = torch.cuda.mem_get_info(device_id)
+
+        if distributed:
+            import torch.distributed as dist
+
+            tensor = torch.tensor(free_gpu_memory, dtype=torch.float32, device="cuda")
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group)
+            free_gpu_memory = float(tensor.item())
+
+        return free_gpu_memory / (1 << 30)
+
     @classmethod
     def get_attn_backend_cls_str(
         cls,
@@ -409,6 +442,7 @@ class NonNvmlCudaPlatform(CudaPlatformBase):
         return str(torch.cuda.get_device_name(device_id))
 
     @classmethod
+    @lru_cache(maxsize=1)
     def get_device_total_memory(cls, device_id: int = 0) -> int:
         device_props = torch.cuda.get_device_properties(device_id)
         return int(device_props.total_memory)

python/sglang/multimodal_gen/runtime/platforms/interface.py

@@ -7,7 +7,7 @@ from __future__ import annotations
 import enum
 import random
 from functools import lru_cache
-from typing import TYPE_CHECKING, NamedTuple
+from typing import TYPE_CHECKING, Any, NamedTuple
 
 import numpy as np
 import torch
@@ -216,6 +216,7 @@ class Platform:
         raise NotImplementedError
 
     @classmethod
+    @lru_cache(maxsize=1)
     def get_device_total_memory(cls, device_id: int = 0) -> int:
         """Get the total memory of a device in bytes."""
         raise NotImplementedError
@@ -307,6 +308,19 @@ class Platform:
         """
         raise NotImplementedError
 
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int = 0,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+        """
+        Return the available memory in GiB.
+        """
+        raise NotImplementedError
+
     @classmethod
     def get_device_communicator_cls(cls) -> str:
         """

python/sglang/multimodal_gen/runtime/platforms/mps.py

@@ -1,17 +1,21 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+from functools import lru_cache
+from typing import Any
 
-# SPDX-License-Identifier: Apache-2.0
-
+import psutil
 import torch
 
-from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
-from sglang.multimodal_gen.runtime.platforms.interface import (
-    DeviceCapability,
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
     Platform,
     PlatformEnum,
 )
+from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
+# SPDX-License-Identifier: Apache-2.0
+
+
 logger = init_logger(__name__)
 
 
@@ -35,8 +39,10 @@ class MpsPlatform(Platform):
         raise NotImplementedError
 
     @classmethod
+    @lru_cache(maxsize=1)
     def get_device_total_memory(cls, device_id: int = 0) -> int:
-        raise NotImplementedError
+
+        return psutil.virtual_memory().total
 
     @classmethod
     def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
@@ -55,6 +61,30 @@ class MpsPlatform(Platform):
     ) -> float:
         return 0.0
 
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int = 0,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+
+        if empty_cache:
+            torch.mps.empty_cache()
+
+        # For MPS, available memory is essentially the system available memory
+        free_memory = psutil.virtual_memory().available
+
+        if distributed:
+            import torch.distributed as dist
+
+            tensor = torch.tensor(free_memory, dtype=torch.float32)
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group)
+            free_memory = float(tensor.item())
+
+        return free_memory / (1 << 30)
+
     @classmethod
     def get_attn_backend_cls_str(
         cls,

python/sglang/multimodal_gen/runtime/platforms/rocm.py

@@ -6,6 +6,8 @@
 This file is a platform abstraction for ROCm GPUs,
 adjusted to match the structure and interface of `cuda.py`.
 """
+from functools import lru_cache
+from typing import Any
 
 import torch
 
@@ -39,6 +41,7 @@ class RocmPlatform(Platform):
         return str(torch.cuda.get_device_name(device_id))
 
     @classmethod
+    @lru_cache(maxsize=1)
     def get_device_total_memory(cls, device_id: int = 0) -> int:
         return torch.cuda.get_device_properties(device_id).total_memory
 
@@ -61,6 +64,28 @@ class RocmPlatform(Platform):
         torch.cuda.reset_peak_memory_stats(device)
         return float(torch.cuda.max_memory_allocated(device))
 
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int = 0,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+        if empty_cache:
+            torch.cuda.empty_cache()
+
+        free_gpu_memory, _ = torch.cuda.mem_get_info(device_id)
+
+        if distributed:
+            import torch.distributed as dist
+
+            tensor = torch.tensor(free_gpu_memory, dtype=torch.float32, device="cuda")
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group)
+            free_gpu_memory = float(tensor.item())
+
+        return free_gpu_memory / (1 << 30)
+
     @classmethod
     def get_attn_backend_cls_str(
         cls,