Tiny call cudaProfilerStart only on first rank in node (#14211)

This commit is contained in:
fzyzcjy
2025-12-01 18:18:45 +08:00
committed by GitHub
parent fa9021b21f
commit d815d00248
2 changed files with 16 additions and 6 deletions

View File

@@ -9,6 +9,7 @@ import torch
from sglang.srt.environ import envs
from sglang.srt.managers.io_struct import ProfileReq, ProfileReqOutput, ProfileReqType
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import is_npu
from sglang.srt.utils.profile_merger import ProfileMerger
from sglang.srt.utils.profile_utils import ProfileManager
@@ -201,7 +202,8 @@ class SchedulerProfilerMixin:
self.profile_in_progress = True
if "CUDA_PROFILER" in activities:
torch.cuda.cudart().cudaProfilerStart()
if self.gpu_id == get_global_server_args().base_gpu_id:
torch.cuda.cudart().cudaProfilerStart()
self.profile_in_progress = True
return ProfileReqOutput(success=True, message="Succeeded")
@@ -310,7 +312,8 @@ class SchedulerProfilerMixin:
torch.cuda.memory._record_memory_history(enabled=None)
if "CUDA_PROFILER" in self.profiler_activities:
torch.cuda.cudart().cudaProfilerStop()
if self.gpu_id == get_global_server_args().base_gpu_id:
torch.cuda.cudart().cudaProfilerStop()
merge_message = self._merge_profile_traces()