chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

7
python/sglang/srt/managers/scheduler_profiler_mixin.py
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()

15
python/sglang/srt/utils/profile_utils.py
View File

@@ -10,6 +10,7 @@ import torch
from sglang.srt.managers.io_struct import ProfileReqOutput
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
_is_npu = is_npu()
@@ -34,6 +35,7 @@ class ProfileManager:
)
self.tp_rank = tp_rank
self.cpu_group = cpu_group
self.first_rank_in_node = gpu_id == get_global_server_args().base_gpu_id
self.profiler_kwargs = None
self.profiler = None
@@ -105,6 +107,7 @@ class ProfileManager:
**self.profiler_kwargs,
tp_rank=self.tp_rank,
cpu_group=self.cpu_group,
first_rank_in_node=self.first_rank_in_node,
output_suffix=f"-{stage}" if stage else "",
)
self.profiler.start()
@@ -239,6 +242,7 @@ class _ProfilerConcreteBase(_ProfilerBase):
profile_id: str,
tp_rank: int,
cpu_group,
first_rank_in_node: bool,
):
self.output_dir = output_dir
self.output_prefix = output_prefix
@@ -246,6 +250,7 @@ class _ProfilerConcreteBase(_ProfilerBase):
self.profile_id = profile_id
self.tp_rank = tp_rank
self.cpu_group = cpu_group
self.first_rank_in_node = first_rank_in_node
class _ProfilerTorch(_ProfilerConcreteBase):
@@ -329,12 +334,14 @@ class _ProfilerMemory(_ProfilerConcreteBase):
class _ProfilerCudart(_ProfilerConcreteBase):
def start(self):
logger.info(f"Call cudaProfilerStart")
torch.cuda.cudart().cudaProfilerStart()
if self.first_rank_in_node:
logger.info(f"Call cudaProfilerStart")
torch.cuda.cudart().cudaProfilerStart()
def stop(self):
logger.info(f"Call cudaProfilerStop")
torch.cuda.cudart().cudaProfilerStop()
if self.first_rank_in_node:
logger.info(f"Call cudaProfilerStop")
torch.cuda.cudart().cudaProfilerStop()
class _ProfilerRPD(_ProfilerConcreteBase):