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

[diffusion] logging: improve peak vram logging (#18865)

Browse Source
This commit is contained in:
Mick
2026-02-16 16:44:37 +08:00
committed by GitHub
parent ed22720c07
commit d0c94e136a
4 changed files with 133 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py
View File

@@ -21,8 +21,9 @@ from sglang.multimodal_gen.runtime.entrypoints.cli.utils import (
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.perf_logger import (
MemorySnapshot,
PerformanceLogger,
RequestTimings,
RequestMetrics,
)
from sglang.multimodal_gen.utils import FlexibleArgumentParser
@@ -72,11 +73,22 @@ def maybe_dump_performance(args: argparse.Namespace, server_args, prompt: str, r
if not (args.perf_dump_path and timings_dict):
return
timings = RequestTimings(request_id=timings_dict.get("request_id"))
timings = RequestMetrics(request_id=timings_dict.get("request_id"))
timings.stages = timings_dict.get("stages", {})
timings.steps = timings_dict.get("steps", [])
timings.total_duration_ms = timings_dict.get("total_duration_ms", 0)
# restore memory snapshots from serialized dict
memory_snapshots_dict = timings_dict.get("memory_snapshots", {})
for checkpoint_name, snapshot_dict in memory_snapshots_dict.items():
snapshot = MemorySnapshot(
allocated_mb=snapshot_dict.get("allocated_mb", 0.0),
reserved_mb=snapshot_dict.get("reserved_mb", 0.0),
peak_allocated_mb=snapshot_dict.get("peak_allocated_mb", 0.0),
peak_reserved_mb=snapshot_dict.get("peak_reserved_mb", 0.0),
)
timings.memory_snapshots[checkpoint_name] = snapshot
PerformanceLogger.dump_benchmark_report(
file_path=args.perf_dump_path,
timings=timings,

41
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
View File

@@ -45,7 +45,10 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import (
globally_suppress_loggers,
init_logger,
)
from sglang.multimodal_gen.runtime.utils.perf_logger import PerformanceLogger
from sglang.multimodal_gen.runtime.utils.perf_logger import (
PerformanceLogger,
capture_memory_snapshot,
)
logger = init_logger(__name__)
@@ -146,11 +149,20 @@ class GPUWorker:
)
def do_mem_analysis(self, output_batch: OutputBatch):
peak_memory_bytes = torch.cuda.max_memory_allocated()
output_batch.peak_memory_mb = peak_memory_bytes / (1024**2)
peak_memory_gb = peak_memory_bytes / (1024**3)
final_snapshot = capture_memory_snapshot()
if output_batch.timings:
output_batch.timings.record_memory_snapshot("mem_analysis", final_snapshot)
# for details on max_memory_reserved: https://docs.pytorch.org/docs/stable/generated/torch.cuda.memory.max_memory_reserved.html
peak_reserved_bytes = torch.cuda.max_memory_reserved()
peak_allocated_bytes = torch.cuda.max_memory_allocated()
output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2)
peak_reserved_gb = peak_reserved_bytes / (1024**3)
peak_allocated_gb = peak_allocated_bytes / (1024**3)
remaining_gpu_mem_gb = (
current_platform.get_device_total_memory() / (1024**3) - peak_memory_gb
current_platform.get_device_total_memory() / (1024**3) - peak_reserved_gb
)
can_stay_resident = self.get_can_stay_resident_components(remaining_gpu_mem_gb)
suggested_args = set()
@@ -173,8 +185,13 @@ class GPUWorker:
suggested_args_str = (
", ".join(sorted(suggested_args)) if suggested_args else "None"
)
pool_overhead_gb = peak_reserved_gb - peak_allocated_gb
logger.info(
f"Peak GPU memory: {peak_memory_gb:.2f} GB, "
f"Peak GPU memory: {peak_reserved_gb:.2f} GB, "
f"Peak allocated: {peak_allocated_gb:.2f} GB, "
f"Memory pool overhead: {pool_overhead_gb:.2f} GB ({pool_overhead_gb/peak_reserved_gb*100:.1f}%), "
f"Remaining GPU memory at peak: {remaining_gpu_mem_gb:.2f} GB. "
f"Components that could stay resident (based on the last request workload): {can_stay_resident}. "
f"Related offload server args to disable: {suggested_args_str}"
@@ -193,6 +210,11 @@ class GPUWorker:
start_time = time.monotonic()
# capture memory baseline before forward
if self.rank == 0 and req.timings:
baseline_snapshot = capture_memory_snapshot()
req.timings.record_memory_snapshot("before_forward", baseline_snapshot)
req.log(server_args=self.server_args)
result = self.pipeline.forward(req, self.server_args)
@@ -210,6 +232,13 @@ class GPUWorker:
else:
output_batch = result
# capture memory after forward (peak)
if self.rank == 0 and output_batch.timings:
peak_snapshot = capture_memory_snapshot()
output_batch.timings.record_memory_snapshot(
"after_forward", peak_snapshot
)
if self.rank == 0 and not req.suppress_logs:
self.do_mem_analysis(output_batch)

10
python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py
View File

@@ -26,7 +26,7 @@ from sglang.multimodal_gen.configs.sample.teacache import (
)
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestTimings
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestMetrics
from sglang.multimodal_gen.utils import align_to
logger = init_logger(__name__)
@@ -152,7 +152,7 @@ class Req:
VSA_sparsity: float = 0.0
# stage logging
timings: Optional["RequestTimings"] = None
timings: Optional["RequestMetrics"] = None
# results
output: torch.Tensor | None = None
@@ -267,7 +267,7 @@ class Req:
if self.guidance_scale_2 is None:
self.guidance_scale_2 = self.guidance_scale
self.timings = RequestTimings(request_id=self.request_id)
self.timings = RequestMetrics(request_id=self.request_id)
if self.is_warmup:
self.set_as_warmup()
@@ -329,8 +329,8 @@ class OutputBatch:
error: str | None = None
output_file_paths: list[str] | None = None
# logged timings info, directly from Req.timings
timings: Optional["RequestTimings"] = None
# logged metrics info, directly from Req.timings
timings: Optional["RequestMetrics"] = None
# For ComfyUI integration: noise prediction from denoising stage
noise_pred: torch.Tensor | None = None

87
python/sglang/multimodal_gen/runtime/utils/perf_logger.py
View File

@@ -25,14 +25,32 @@ logger = init_logger(__name__)
@dataclasses.dataclass
class RequestTimings:
"""A lightweight data class to store performance timings for a single request."""
class MemorySnapshot:
allocated_mb: float # current allocated memory
reserved_mb: float # current reserved memory (actual VRAM)
peak_allocated_mb: float # peak allocated since last reset
peak_reserved_mb: float # peak reserved since last reset
def to_dict(self) -> Dict[str, Any]:
return {
"allocated_mb": round(self.allocated_mb, 2),
"reserved_mb": round(self.reserved_mb, 2),
"peak_allocated_mb": round(self.peak_allocated_mb, 2),
"peak_reserved_mb": round(self.peak_reserved_mb, 2),
}
@dataclasses.dataclass
class RequestMetrics:
"""Performance metrics for a single request, including timings and memory snapshots."""
def __init__(self, request_id: str):
self.request_id = request_id
self.stages: Dict[str, float] = {}
self.steps: list[float] = []
self.total_duration_ms: float = 0.0
# memory tracking: {checkpoint_name: MemorySnapshot}
self.memory_snapshots: Dict[str, MemorySnapshot] = {}
@property
def total_duration_s(self) -> float:
@@ -47,13 +65,20 @@ class RequestTimings:
assert index == len(self.steps)
self.steps.append(duration_s * 1000)
def record_memory_snapshot(self, checkpoint_name: str, snapshot: MemorySnapshot):
self.memory_snapshots[checkpoint_name] = snapshot
def to_dict(self) -> Dict[str, Any]:
"""Serializes the timing data to a dictionary."""
"""Serializes the metrics data to a dictionary."""
return {
"request_id": self.request_id,
"stages": self.stages,
"steps": self.steps,
"total_duration_ms": self.total_duration_ms,
"memory_snapshots": {
name: snapshot.to_dict()
for name, snapshot in self.memory_snapshots.items()
},
}
@@ -90,6 +115,28 @@ def get_git_commit_hash() -> str:
return "N/A"
def capture_memory_snapshot() -> MemorySnapshot:
if not torch.cuda.is_available():
return MemorySnapshot(
allocated_mb=0.0,
reserved_mb=0.0,
peak_allocated_mb=0.0,
peak_reserved_mb=0.0,
)
allocated = torch.cuda.memory_allocated()
reserved = torch.cuda.memory_reserved()
peak_allocated = torch.cuda.max_memory_allocated()
peak_reserved = torch.cuda.max_memory_reserved()
return MemorySnapshot(
allocated_mb=allocated / (1024**2),
reserved_mb=reserved / (1024**2),
peak_allocated_mb=peak_allocated / (1024**2),
peak_reserved_mb=peak_reserved / (1024**2),
)
@dataclasses.dataclass
class RequestPerfRecord:
request_id: str
@@ -101,6 +148,7 @@ class RequestPerfRecord:
stages: list[dict]
steps: list[float]
total_duration_ms: float
memory_snapshots: dict[str, dict] = dataclasses.field(default_factory=dict)
def __init__(
self,
@@ -110,6 +158,7 @@ class RequestPerfRecord:
stages,
steps,
total_duration_ms,
memory_snapshots=None,
timestamp=None,
):
self.request_id = request_id
@@ -123,20 +172,22 @@ class RequestPerfRecord:
self.stages = stages
self.steps = steps
self.total_duration_ms = total_duration_ms
self.memory_snapshots = memory_snapshots or {}
class StageProfiler:
"""
A unified context manager, records timing information (usually of a single Stage or a step) into a provided RequestTimings object (usually from a Req).
A unified context manager, records performance metrics (usually of a single Stage or a step) into a provided RequestMetrics object (usually from a Req).
"""
def __init__(
self,
stage_name: str,
logger: _SGLDiffusionLogger,
timings: Optional["RequestTimings"],
timings: Optional["RequestMetrics"],
log_stage_start_end: bool = False,
perf_dump_path_provided: bool = False,
capture_memory: bool = False,
):
self.stage_name = stage_name
self.timings = timings
@@ -144,6 +195,7 @@ class StageProfiler:
self.start_time = 0.0
self.log_timing = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING
self.log_stage_start_end = log_stage_start_end
self.capture_memory = capture_memory
def __enter__(self):
if self.log_stage_start_end:
@@ -194,6 +246,13 @@ class StageProfiler:
else:
self.timings.record_stage(self.stage_name, execution_time_s)
# capture memory snapshot after stage if requested
if self.capture_memory and torch.cuda.is_available():
snapshot = capture_memory_snapshot()
self.timings.record_memory_snapshot(
f"after_{self.stage_name}", snapshot
)
return False
@@ -203,14 +262,14 @@ class PerformanceLogger:
Serves both as a runtime logger (stream to file) and a dump utility.
Notice that ""RequestTimings"" stores the performance metrics of a single request
Notice that RequestMetrics stores the performance metrics of a single request
"""
@classmethod
def dump_benchmark_report(
cls,
file_path: str,
timings: "RequestTimings",
timings: "RequestMetrics",
meta: Optional[Dict[str, Any]] = None,
tag: str = "benchmark_dump",
):
@@ -228,6 +287,11 @@ class PerformanceLogger:
for idx, duration_ms in enumerate(timings.steps)
]
memory_checkpoints = {
name: snapshot.to_dict()
for name, snapshot in timings.memory_snapshots.items()
}
report = {
"timestamp": datetime.now(UTC).isoformat(),
"request_id": timings.request_id,
@@ -236,6 +300,7 @@ class PerformanceLogger:
"total_duration_ms": timings.total_duration_ms,
"steps": formatted_steps,
"denoise_steps_ms": denoise_steps_ms,
"memory_checkpoints": memory_checkpoints,
"meta": meta or {},
}
@@ -251,7 +316,7 @@ class PerformanceLogger:
@classmethod
def log_request_summary(
cls,
timings: "RequestTimings",
timings: "RequestMetrics",
tag: str = "total_inference_time",
):
"""logs the stage metrics and total duration for a completed request
@@ -264,6 +329,11 @@ class PerformanceLogger:
for name, duration_ms in timings.stages.items()
]
memory_checkpoints = {
name: snapshot.to_dict()
for name, snapshot in timings.memory_snapshots.items()
}
record = RequestPerfRecord(
timings.request_id,
commit_hash=get_git_commit_hash(),
@@ -271,6 +341,7 @@ class PerformanceLogger:
stages=formatted_stages,
steps=timings.steps,
total_duration_ms=timings.total_duration_ms,
memory_snapshots=memory_checkpoints,
)
try: