505 lines
17 KiB
Python
505 lines
17 KiB
Python
import logging
|
|
import os
|
|
import time
|
|
from abc import ABC
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Callable, Dict, List, Optional
|
|
|
|
import torch
|
|
|
|
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.managers.io_struct import ProfileReqOutput
|
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
|
from sglang.srt.model_executor.step_span_utils import (
|
|
build_detailed_annotation_suffix,
|
|
detailed_annotations_enabled,
|
|
set_detailed_annotations_enabled,
|
|
)
|
|
from sglang.srt.platforms import current_platform
|
|
from sglang.srt.runtime_context import get_device, get_parallel
|
|
from sglang.srt.utils import is_npu
|
|
from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches
|
|
|
|
_is_npu = is_npu()
|
|
if _is_npu:
|
|
import torch_npu
|
|
|
|
patches = [
|
|
["profiler.profile", torch_npu.profiler.profile],
|
|
["profiler.ProfilerActivity.CUDA", torch_npu.profiler.ProfilerActivity.NPU],
|
|
["profiler.ProfilerActivity.CPU", torch_npu.profiler.ProfilerActivity.CPU],
|
|
]
|
|
apply_torch_npu_patches(torch_npu, patches)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Single source of truth for the CUDA-graph capture trace output directory,
|
|
# shared by both the original single-trace export and the per-batch-size
|
|
# (SGLANG_GRAPH_BATCH_CAPTURE) traces so they land in the same place.
|
|
GRAPH_CAPTURE_PROFILE_DIRNAME = "graph_capture_profile"
|
|
|
|
|
|
def graph_capture_profile_dir() -> str:
|
|
"""``<SGLANG_TORCH_PROFILER_DIR>/graph_capture_profile`` — the one directory
|
|
both capture-trace modes write to. Change the location here only."""
|
|
return os.path.join(
|
|
envs.SGLANG_TORCH_PROFILER_DIR.get(), GRAPH_CAPTURE_PROFILE_DIRNAME
|
|
)
|
|
|
|
|
|
def export_cuda_graph_capture_trace(prof_context, *, runner_name: str, tp_rank: int):
|
|
"""Persist a CUDA-graph capture profiler trace (chrome trace) to disk.
|
|
|
|
Opt-in via ``SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE`` (no-op otherwise). The
|
|
capture profiler must have run with ``record_shapes=True`` so the trace can
|
|
be inspected offline as a per-kernel shape/identity record. The file lands in
|
|
``graph_capture_profile_dir()`` and is namespaced by runner class and TP rank
|
|
so concurrent capture passes (e.g. EAGLE3 target/draft/draft-extend) and
|
|
ranks don't overwrite each other.
|
|
"""
|
|
if not envs.SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE.get():
|
|
return
|
|
output_dir = graph_capture_profile_dir()
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
path = os.path.join(
|
|
output_dir, f"cuda_graph_capture-{runner_name}-TP-{tp_rank}.json.gz"
|
|
)
|
|
prof_context.export_chrome_trace(path)
|
|
logger.info(f"CUDA graph capture trace saved to: {path}")
|
|
|
|
|
|
class ProfileManager:
|
|
def __init__(self, ps: ParallelState, cpu_group):
|
|
self.stage_based_trigger = _StageBasedTrigger(
|
|
on_start=self._do_start,
|
|
on_stop=self._do_stop,
|
|
)
|
|
self.ps = ps
|
|
self.cpu_group = cpu_group
|
|
self.first_rank_in_node = ps.gpu_id == get_device().base_gpu_id
|
|
self.profiler_kwargs = None
|
|
self.profiler = None
|
|
self.detailed_annotations = False
|
|
|
|
def step(self, forward_mode: ForwardMode):
|
|
stage = _get_stage_from_forward_mode(forward_mode)
|
|
if stage is None:
|
|
return
|
|
|
|
self.stage_based_trigger.step(stage=stage)
|
|
|
|
def configure(
|
|
self,
|
|
*,
|
|
output_dir: Optional[str],
|
|
start_step: Optional[int],
|
|
num_steps: Optional[int],
|
|
activities: Optional[List[str]],
|
|
with_stack: Optional[bool],
|
|
record_shapes: Optional[bool],
|
|
profile_by_stage: bool,
|
|
profile_id: str,
|
|
merge_profiles: bool,
|
|
profile_prefix: str,
|
|
profile_stages: Optional[List[str]] = None,
|
|
detailed_annotations: bool = False,
|
|
):
|
|
self.detailed_annotations = detailed_annotations
|
|
# not supported yet
|
|
assert start_step is None
|
|
assert profile_by_stage, (
|
|
"only support profile_by_stage=true now"
|
|
) # `false` can be easily supported
|
|
assert not merge_profiles
|
|
|
|
if output_dir is None:
|
|
output_dir = os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp")
|
|
if activities is None:
|
|
activities = ["CPU", "GPU"]
|
|
|
|
self.profiler_kwargs = dict(
|
|
activities=activities,
|
|
with_stack=with_stack,
|
|
record_shapes=record_shapes,
|
|
output_dir=output_dir,
|
|
output_prefix=profile_prefix,
|
|
profile_id=profile_id,
|
|
)
|
|
|
|
self.stage_based_trigger.configure(
|
|
num_steps=num_steps,
|
|
interesting_stages=profile_stages or ["prefill", "decode"],
|
|
)
|
|
|
|
return ProfileReqOutput(success=True, message="Succeeded")
|
|
|
|
def manual_start(self):
|
|
raise NotImplementedError("manually start is only supported yet")
|
|
|
|
def manual_stop(self):
|
|
raise NotImplementedError("manually stop is only supported yet")
|
|
|
|
def _do_start(self, stage: Optional[str] = None):
|
|
logger.info(
|
|
f"Profiling starts{f' for {stage}' if stage else ''}. "
|
|
f"Traces will be saved to: {self.profiler_kwargs['output_dir']} "
|
|
f"(with profile id: {self.profiler_kwargs['profile_id']})",
|
|
)
|
|
|
|
assert self.profiler is None
|
|
# Fold the per-phase c_/g_ aggregates into the step span while this
|
|
# stage's profile is active (v2 auto-start path; reset in _do_stop).
|
|
set_detailed_annotations_enabled(self.detailed_annotations)
|
|
self.profiler = _ProfilerBase.create(
|
|
**self.profiler_kwargs,
|
|
ps=self.ps,
|
|
cpu_group=self.cpu_group,
|
|
first_rank_in_node=self.first_rank_in_node,
|
|
output_suffix=f"-{stage}" if stage else "",
|
|
)
|
|
self.profiler.start()
|
|
|
|
def _do_stop(self):
|
|
logger.info("Stop profiling...")
|
|
self.profiler.stop()
|
|
logger.info(
|
|
f"Profiling done. Traces are saved to: {self.profiler_kwargs['output_dir']}"
|
|
)
|
|
self.profiler = None
|
|
# Clear the detailed step-span toggle here too: the v2 trigger auto-stop
|
|
# goes through _do_stop (not SchedulerProfilerManager._stop_profile), so
|
|
# this guarantees the flag resets on every stop path.
|
|
set_detailed_annotations_enabled(False)
|
|
|
|
|
|
def _get_stage_from_forward_mode(forward_mode: ForwardMode):
|
|
if forward_mode.is_prefill():
|
|
return "prefill"
|
|
elif forward_mode.is_decode():
|
|
return "decode"
|
|
elif forward_mode.is_idle():
|
|
return None
|
|
else:
|
|
raise RuntimeError(f"unsupported profile stage: {forward_mode=}")
|
|
|
|
|
|
# ======================================== Stage related ==========================================
|
|
|
|
|
|
class _StageBasedTrigger:
|
|
@dataclass
|
|
class _StageConfig:
|
|
target_count: int
|
|
|
|
@dataclass
|
|
class _RunningState:
|
|
curr_stage: str
|
|
curr_count: int
|
|
|
|
def __init__(self, on_start: Callable, on_stop: Callable):
|
|
self.on_start = on_start
|
|
self.on_stop = on_stop
|
|
|
|
self.running_state: Optional[_StageBasedTrigger._RunningState] = None
|
|
# When a stage is in the dict, it means it is being or should be executed
|
|
self.stage_configs: Dict[str, _StageBasedTrigger._StageConfig] = {}
|
|
|
|
def configure(self, num_steps: int, interesting_stages: List[str]):
|
|
assert self.running_state is None
|
|
self.stage_configs = {
|
|
stage: self._StageConfig(target_count=num_steps)
|
|
for stage in interesting_stages
|
|
}
|
|
|
|
def step(self, stage: str):
|
|
# Incr counter
|
|
if (s := self.running_state) is not None:
|
|
s.curr_count += 1
|
|
|
|
# Maybe stop
|
|
if ((s := self.running_state) is not None) and (
|
|
(s.curr_count > self.stage_configs[s.curr_stage].target_count)
|
|
or (stage != s.curr_stage)
|
|
):
|
|
del self.stage_configs[s.curr_stage]
|
|
self.running_state = None
|
|
self.on_stop()
|
|
|
|
# Maybe start
|
|
if (self.running_state is None) and (stage in self.stage_configs):
|
|
self.running_state = self._RunningState(
|
|
curr_stage=stage,
|
|
curr_count=0,
|
|
)
|
|
self.on_start(stage=stage)
|
|
|
|
# Sanity check
|
|
assert (self.running_state is not None) == (stage in self.stage_configs)
|
|
if (s := self.running_state) is not None:
|
|
assert s.curr_stage == stage
|
|
|
|
|
|
# ======================================== Concrete profilers ==========================================
|
|
|
|
|
|
class _ProfilerBase(ABC):
|
|
@staticmethod
|
|
def create(activities, with_stack, record_shapes, **kwargs):
|
|
inners = []
|
|
if current_platform.is_out_of_tree():
|
|
if current_platform.get_torch_profiler_activity_str() in activities:
|
|
inners.append(
|
|
_ProfilerTorch(
|
|
**kwargs,
|
|
activities=activities,
|
|
with_stack=with_stack,
|
|
record_shapes=record_shapes,
|
|
)
|
|
)
|
|
if ("CPU" in activities) or ("GPU" in activities):
|
|
inners.append(
|
|
_ProfilerTorch(
|
|
**kwargs,
|
|
activities=activities,
|
|
with_stack=with_stack,
|
|
record_shapes=record_shapes,
|
|
)
|
|
)
|
|
if "MEM" in activities:
|
|
inners.append(_ProfilerMemory(**kwargs))
|
|
if "CUDA_PROFILER" in activities:
|
|
inners.append(_ProfilerCudart(**kwargs))
|
|
if "RPD" in activities: # for ROCM
|
|
inners.append(_ProfilerRPD(**kwargs))
|
|
|
|
return _ProfilerList(inners)
|
|
|
|
def start(self):
|
|
raise NotImplementedError
|
|
|
|
def stop(self):
|
|
raise NotImplementedError
|
|
|
|
|
|
class _ProfilerList(_ProfilerBase):
|
|
def __init__(self, inners: List[_ProfilerBase]):
|
|
self.inners = inners
|
|
|
|
def start(self):
|
|
for inner in self.inners:
|
|
inner.start()
|
|
|
|
def stop(self):
|
|
for inner in self.inners:
|
|
inner.stop()
|
|
|
|
|
|
class _ProfilerConcreteBase(_ProfilerBase):
|
|
def __init__(
|
|
self,
|
|
output_dir: str,
|
|
output_prefix: str,
|
|
output_suffix: str,
|
|
profile_id: str,
|
|
ps: ParallelState,
|
|
cpu_group,
|
|
first_rank_in_node: bool,
|
|
):
|
|
self.output_dir = output_dir
|
|
self.output_prefix = output_prefix
|
|
self.output_suffix = output_suffix
|
|
self.profile_id = profile_id
|
|
self.ps = ps
|
|
self.cpu_group = cpu_group
|
|
self.first_rank_in_node = first_rank_in_node
|
|
|
|
|
|
class _ProfilerTorch(_ProfilerConcreteBase):
|
|
def __init__(self, with_stack: bool, record_shapes: bool, activities, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.with_stack = with_stack
|
|
self.record_shapes = record_shapes
|
|
self.activities = activities
|
|
|
|
def start(self):
|
|
activity_map = {
|
|
"CPU": torch.profiler.ProfilerActivity.CPU,
|
|
"GPU": torch.profiler.ProfilerActivity.CUDA,
|
|
}
|
|
|
|
if current_platform.is_out_of_tree():
|
|
activity_map[current_platform.get_torch_profiler_activity_str()] = (
|
|
current_platform.get_torch_profiler_activity()
|
|
)
|
|
|
|
torchprof_activities = [
|
|
activity_map[a] for a in self.activities if a in activity_map
|
|
]
|
|
|
|
self.torch_profiler = torch.profiler.profile(
|
|
activities=torchprof_activities,
|
|
with_stack=self.with_stack if self.with_stack is not None else True,
|
|
record_shapes=(
|
|
self.record_shapes if self.record_shapes is not None else False
|
|
),
|
|
on_trace_ready=(
|
|
None
|
|
if not _is_npu
|
|
else torch_npu.profiler.tensorboard_trace_handler(self.output_dir)
|
|
),
|
|
)
|
|
self.torch_profiler.start()
|
|
|
|
def stop(self):
|
|
Path(self.output_dir).mkdir(parents=True, exist_ok=True)
|
|
|
|
self.torch_profiler.stop()
|
|
if not _is_npu:
|
|
# Build filename with only non-zero ranks to maintain backward compatibility
|
|
filename_parts = [self.profile_id, f"TP-{get_parallel().tp_rank}"]
|
|
|
|
# Only add other ranks if parallelism is enabled (size > 1)
|
|
if self.ps.dp_size > 1:
|
|
filename_parts.append(f"DP-{get_parallel().dp_rank}")
|
|
if get_parallel().pp_size > 1:
|
|
filename_parts.append(f"PP-{get_parallel().pp_rank}")
|
|
if get_parallel().moe_ep_size > 1:
|
|
filename_parts.append(f"EP-{get_parallel().moe_ep_rank}")
|
|
|
|
filename = (
|
|
(self.output_prefix + "-" if self.output_prefix else "")
|
|
+ "-".join(filename_parts)
|
|
+ self.output_suffix
|
|
+ ".trace.json.gz"
|
|
)
|
|
|
|
self.torch_profiler.export_chrome_trace(
|
|
os.path.join(self.output_dir, filename)
|
|
)
|
|
torch.distributed.barrier(self.cpu_group)
|
|
|
|
# TODO: migrate `_merge_profile_traces`
|
|
|
|
|
|
class _ProfilerMemory(_ProfilerConcreteBase):
|
|
def start(self):
|
|
torch.cuda.memory._record_memory_history(
|
|
max_entries=envs.SGLANG_MEM_PROFILE_MAX_ENTRIES.get()
|
|
)
|
|
|
|
def stop(self):
|
|
Path(self.output_dir).mkdir(parents=True, exist_ok=True)
|
|
|
|
memory_profile_path = os.path.join(
|
|
self.output_dir,
|
|
(self.output_prefix + "-" if self.output_prefix else "")
|
|
+ str(time.time())
|
|
+ f"-TP-{get_parallel().tp_rank}-memory"
|
|
+ self.output_suffix
|
|
+ ".pickle",
|
|
)
|
|
torch.cuda.memory._dump_snapshot(memory_profile_path)
|
|
torch.cuda.memory._record_memory_history(enabled=None)
|
|
|
|
|
|
class _ProfilerCudart(_ProfilerConcreteBase):
|
|
def start(self):
|
|
if self.first_rank_in_node:
|
|
logger.info(f"Call cudaProfilerStart")
|
|
torch.cuda.cudart().cudaProfilerStart()
|
|
|
|
def stop(self):
|
|
if self.first_rank_in_node:
|
|
logger.info(f"Call cudaProfilerStop")
|
|
torch.cuda.cudart().cudaProfilerStop()
|
|
|
|
|
|
class _ProfilerRPD(_ProfilerConcreteBase):
|
|
def start(self):
|
|
Path(self.output_dir).mkdir(parents=True, exist_ok=True)
|
|
|
|
from rpdTracerControl import rpdTracerControl
|
|
|
|
rpdTracerControl.skipCreate()
|
|
|
|
self.rpd_profile_path = os.path.join(
|
|
self.output_dir,
|
|
"rpd-"
|
|
+ str(time.time())
|
|
+ f"-TP-{get_parallel().tp_rank}"
|
|
+ ".trace.json.gz",
|
|
)
|
|
|
|
if get_parallel().tp_rank == 0:
|
|
import sqlite3
|
|
|
|
from rocpd.schema import RocpdSchema
|
|
|
|
if os.path.exists("trace.rpd"):
|
|
os.unlink("trace.rpd")
|
|
schema = RocpdSchema()
|
|
connection = sqlite3.connect("trace.rpd")
|
|
schema.writeSchema(connection)
|
|
connection.commit()
|
|
del connection
|
|
torch.distributed.barrier(self.cpu_group)
|
|
|
|
self.rpd_profiler = rpdTracerControl()
|
|
self.rpd_profiler.setPythonTrace(True)
|
|
self.rpd_profiler.start()
|
|
self.rpd_profiler.rangePush("", "rpd profile range", "")
|
|
|
|
def stop(self):
|
|
self.rpd_profiler.rangePop()
|
|
self.rpd_profiler.stop()
|
|
self.rpd_profiler.flush()
|
|
|
|
torch.distributed.barrier(self.cpu_group)
|
|
if get_parallel().tp_rank == 0:
|
|
from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace
|
|
|
|
rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path)
|
|
|
|
|
|
def build_step_span_name(
|
|
forward_batch: ForwardBatch,
|
|
detailed_annotations: bool | None = None,
|
|
*,
|
|
is_draft_worker: bool = False,
|
|
) -> str:
|
|
"""Build the profile-trace span name for one forward step.
|
|
|
|
Detailed annotations are folded into the label (via
|
|
build_detailed_annotation_suffix) when enabled. detailed_annotations
|
|
defaults to the process-wide toggle (detailed_annotations_enabled, set
|
|
by the profiler manager); pass an explicit bool to override (e.g. in tests).
|
|
|
|
The target-verify step is labeled ``VERIFY``; every step a draft model
|
|
runner emits is labeled ``DRAFT`` (some draft paths borrow the TARGET_VERIFY
|
|
mode, so the mode name alone cannot tell the two models apart).
|
|
"""
|
|
if detailed_annotations is None:
|
|
detailed_annotations = detailed_annotations_enabled()
|
|
|
|
mode = forward_batch.forward_mode
|
|
bs = forward_batch.batch_size
|
|
if is_draft_worker:
|
|
stage = "DRAFT"
|
|
elif mode == ForwardMode.TARGET_VERIFY:
|
|
stage = "VERIFY"
|
|
else:
|
|
stage = mode.name
|
|
if mode == ForwardMode.EXTEND:
|
|
ext_toks = forward_batch.extend_num_tokens or 0
|
|
base = f"step[{stage} bs={bs} toks={ext_toks}"
|
|
else:
|
|
base = f"step[{stage} bs={bs}"
|
|
|
|
if detailed_annotations:
|
|
suffix = build_detailed_annotation_suffix(forward_batch)
|
|
if suffix:
|
|
base = f"{base} {suffix}"
|
|
return f"{base}]"
|