From 4c403045ec690dbcf3b63a941356e004201ba337 Mon Sep 17 00:00:00 2001
From: wxy <1908865287@qq.com>
Date: Wed, 4 Feb 2026 21:25:23 +0800
Subject: [PATCH] [diffusion] fix: fix the bug of redundant memory usage on
 GPU-0 (#18221)

---
 python/sglang/multimodal_gen/runtime/platforms/cuda.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py
index d59326de1..6b4530a68 100644
--- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py
+++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py
@@ -124,6 +124,9 @@ class CudaPlatformBase(Platform):
         if empty_cache:
             torch.cuda.empty_cache()
 
+        if torch.distributed.is_initialized():
+            device_id = torch.distributed.get_rank()
+
         device_props = torch.cuda.get_device_properties(device_id)
         if device_props.is_integrated:
             free_gpu_memory = psutil.virtual_memory().available