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

[diffusion] logging: log available mem when each stage starts in debug level (#18998)

Browse Source
This commit is contained in:
Mick
2026-02-20 19:57:06 +08:00
committed by GitHub
parent 0d20cf5a66
commit 38a69652e6
19 changed files with 109 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -18,9 +18,11 @@ from sglang.multimodal_gen.runtime.entrypoints.cli.cli_types import CLISubcomman
from sglang.multimodal_gen.runtime.entrypoints.cli.utils import (
RaiseNotImplementedAction,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import GenerationResult
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,
RequestMetrics,
)
@@ -58,7 +60,12 @@ def add_multimodal_gen_generate_args(parser: argparse.ArgumentParser):
return parser
def maybe_dump_performance(args: argparse.Namespace, server_args, prompt: str, results):
def maybe_dump_performance(
args: argparse.Namespace,
server_args,
prompt: str,
results: GenerationResult | list[GenerationResult] | None,
):
"""dump performance if necessary"""
if not (args.perf_dump_path and results):
return
@@ -68,20 +75,29 @@ def maybe_dump_performance(args: argparse.Namespace, server_args, prompt: str, r
else:
result = results
timings_dict = getattr(result, "timings", None) or (
result.get("timings") if isinstance(result, dict) else None
)
if not (args.perf_dump_path and timings_dict):
metrics_dict = result.metrics
if not (args.perf_dump_path and metrics_dict):
return
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)
metrics = RequestMetrics(request_id=metrics_dict.get("request_id"))
metrics.stages = metrics_dict.get("stages", {})
metrics.steps = metrics_dict.get("steps", [])
metrics.total_duration_ms = metrics_dict.get("total_duration_ms", 0)
# restore memory snapshots from serialized dict
memory_snapshots_dict = metrics_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),
)
metrics.memory_snapshots[checkpoint_name] = snapshot
PerformanceLogger.dump_benchmark_report(
file_path=args.perf_dump_path,
timings=timings,
metrics=metrics,
meta={
"prompt": prompt,
"model": server_args.model_path,

8
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
View File

@@ -206,9 +206,9 @@ class DiffGenerator:
size=(req.height, req.width, req.num_frames),
generation_time=timer.duration,
peak_memory_mb=output_batch.peak_memory_mb,
timings=(
output_batch.timings.to_dict()
if output_batch.timings
metrics=(
output_batch.metrics.to_dict()
if output_batch.metrics
else {}
),
trajectory_latents=output_batch.trajectory_latents,
@@ -297,7 +297,7 @@ class DiffGenerator:
if not results:
return
if self.server_args.warmup:
total_duration_ms = results[0].timings.get("total_duration_ms", 0)
total_duration_ms = results[0].metrics.get("total_duration_ms", 0)
logger.info(
f"Warmed-up request processed in {GREEN}%.2f{RESET} seconds (with warmup excluded)",
total_duration_ms / 1000.0,

4
python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py
View File

@@ -304,8 +304,8 @@ def add_common_data_to_response(
if result.peak_memory_mb and result.peak_memory_mb > 0:
response["peak_memory_mb"] = result.peak_memory_mb
if result.timings and result.timings.total_duration_s > 0:
response["inference_time_s"] = result.timings.total_duration_s
if result.metrics and result.metrics.total_duration_s > 0:
response["inference_time_s"] = result.metrics.total_duration_s
response["id"] = request_id

2
python/sglang/multimodal_gen/runtime/entrypoints/utils.py
View File

@@ -105,7 +105,7 @@ class GenerationResult:
size: tuple | None = None # (height, width, num_frames)
generation_time: float = 0.0
peak_memory_mb: float = 0.0
timings: dict = field(default_factory=dict)
metrics: dict = field(default_factory=dict)
trajectory_latents: Any = None
trajectory_timesteps: Any = None
trajectory_decoded: Any = None

7
python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
View File

@@ -19,7 +19,6 @@ from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.loader.utils import (
_normalize_component_type,
component_name_to_loader_cls,
get_memory_usage_of_component,
)
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.server_args import ServerArgs
@@ -125,11 +124,9 @@ class ComponentLoader(ABC):
if isinstance(component, nn.Module):
component = component.eval()
current_gpu_mem = current_platform.get_available_gpu_memory()
consumed = get_memory_usage_of_component(component)
if consumed is None or consumed == 0.0:
consumed = gpu_mem_before_loading - current_gpu_mem
consumed = gpu_mem_before_loading - current_gpu_mem
logger.info(
f"Loaded %s: %s ({source} version). model size: %.2f GB, avail mem: %.2f GB",
f"Loaded %s: %s ({source} version). consumed: %.2f GB, avail mem: %.2f GB",
component_name,
component.__class__.__name__,
consumed,

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

@@ -158,8 +158,8 @@ class GPUWorker:
def do_mem_analysis(self, output_batch: OutputBatch):
final_snapshot = capture_memory_snapshot()
if output_batch.timings:
output_batch.timings.record_memory_snapshot("mem_analysis", final_snapshot)
if output_batch.metrics:
output_batch.metrics.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.get_device_module().max_memory_reserved()
@@ -219,9 +219,9 @@ class GPUWorker:
start_time = time.monotonic()
# capture memory baseline before forward
if self.rank == 0 and req.timings:
if self.rank == 0 and req.metrics:
baseline_snapshot = capture_memory_snapshot()
req.timings.record_memory_snapshot("before_forward", baseline_snapshot)
req.metrics.record_memory_snapshot("before_forward", baseline_snapshot)
req.log(server_args=self.server_args)
result = self.pipeline.forward(req, self.server_args)
@@ -231,7 +231,7 @@ class GPUWorker:
output=result.output,
audio=getattr(result, "audio", None),
audio_sample_rate=getattr(result, "audio_sample_rate", None),
timings=result.timings,
metrics=result.metrics,
trajectory_timesteps=getattr(result, "trajectory_timesteps", None),
trajectory_latents=getattr(result, "trajectory_latents", None),
noise_pred=getattr(result, "noise_pred", None),
@@ -241,9 +241,9 @@ class GPUWorker:
output_batch = result
# capture memory after forward (peak)
if self.rank == 0 and output_batch.timings:
if self.rank == 0 and output_batch.metrics:
peak_snapshot = capture_memory_snapshot()
output_batch.timings.record_memory_snapshot(
output_batch.metrics.record_memory_snapshot(
"after_forward", peak_snapshot
)
@@ -251,7 +251,7 @@ class GPUWorker:
self.do_mem_analysis(output_batch)
duration_ms = (time.monotonic() - start_time) * 1000
output_batch.timings.total_duration_ms = duration_ms
output_batch.metrics.total_duration_ms = duration_ms
# Save output to file and return file path only if requested. Avoid the serialization
# and deserialization overhead between scheduler_client and gpu_worker.
@@ -273,11 +273,13 @@ class GPUWorker:
if req.perf_dump_path is not None or envs.SGLANG_DIFFUSION_STAGE_LOGGING:
# Avoid logging warmup perf records that share the same request_id.
if not req.is_warmup:
PerformanceLogger.log_request_summary(timings=output_batch.timings)
PerformanceLogger.log_request_summary(metrics=output_batch.metrics)
except Exception as e:
logger.error(
f"Error executing request {req.request_id}: {e}", exc_info=True
)
if isinstance(e, _oom_exceptions()):
logger.warning(OOM_MSG)
if output_batch is None:
output_batch = OutputBatch()
output_batch.error = f"Error executing request {req.request_id}: {e}"
@@ -427,8 +429,8 @@ OOM detected. Possible solutions:
- If the OOM occurs during loading:
1. Enable CPU offload for memory-intensive components, or use `--dit-layerwise-offload` for DiT
- If the OOM occurs during runtime:
1. Reduce the number of output tokens by lowering resolution or decreasing `--num-frames`
2. Enable SP and/or TP
1. Enable SP and/or TP (in a multi-GPU setup)
2. Reduce the number of output tokens by lowering resolution or decreasing `--num-frames`
3. Opt for a sparse-attention backend
4. Enable FSDP by `--use-fsdp-inference` (in a multi-GPU setup)
5. Enable quantization (e.g. nunchaku)
@@ -436,6 +438,14 @@ OOM detected. Possible solutions:
"""
def _oom_exceptions():
# torch.OutOfMemoryError exists only in some PyTorch builds
types = [torch.cuda.OutOfMemoryError]
if hasattr(torch, "OutOfMemoryError"):
types.append(torch.OutOfMemoryError)
return tuple(types)
def run_scheduler_process(
local_rank: int,
rank: int,
@@ -483,7 +493,7 @@ def run_scheduler_process(
}
)
scheduler.event_loop()
except torch.OutOfMemoryError as _e:
except _oom_exceptions() as _e:
logger.warning(OOM_MSG)
raise
finally:

4
python/sglang/multimodal_gen/runtime/managers/scheduler.py
View File

@@ -395,12 +395,12 @@ class Scheduler:
if self._warmup_total > 0:
logger.info(
f"Warmup req ({self._warmup_processed}/{self._warmup_total}) processed in {GREEN}%.2f{RESET} seconds",
output_batch.timings.total_duration_s,
output_batch.metrics.total_duration_s,
)
else:
logger.info(
f"Warmup req processed in {GREEN}%.2f{RESET} seconds",
output_batch.timings.total_duration_s,
output_batch.metrics.total_duration_s,
)
else:
if self._warmup_total > 0:

7
python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
View File

@@ -22,6 +22,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor im
)
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
from sglang.multimodal_gen.runtime.pipelines_core.stages import PipelineStage
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
maybe_download_model,
@@ -311,7 +312,11 @@ class ComposedPipelineBase(ABC):
f"Required module: {module_name} was not found in loaded modules: {list(loaded_components.keys())}"
)
logger.debug("Memory usage of loaded modules: %s", self.memory_usages)
logger.debug(
"Memory usage of loaded modules (GiB): %s. Available memory: %s",
self.memory_usages,
round(current_platform.get_available_gpu_memory(), 2),
)
return loaded_components

2
python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py
View File

@@ -31,7 +31,7 @@ class Timer(StageProfiler):
def __init__(self, name="Stage"):
super().__init__(
stage_name=name, timings=None, log_stage_start_end=True, logger=logger
stage_name=name, logger=logger, metrics=None, log_stage_start_end=True
)

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

@@ -152,7 +152,7 @@ class Req:
VSA_sparsity: float = 0.0
# stage logging
timings: Optional["RequestMetrics"] = None
metrics: 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 = RequestMetrics(request_id=self.request_id)
self.metrics = RequestMetrics(request_id=self.request_id)
if self.is_warmup:
self.set_as_warmup()
@@ -330,7 +330,7 @@ class OutputBatch:
output_file_paths: list[str] | None = None
# logged metrics info, directly from Req.timings
timings: Optional["RequestMetrics"] = None
metrics: Optional["RequestMetrics"] = None
# For ComfyUI integration: noise prediction from denoising stage
noise_pred: torch.Tensor | None = None

4
python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py
View File

@@ -195,10 +195,10 @@ class PipelineStage(ABC):
with StageProfiler(
stage_name,
logger=logger,
timings=batch.timings,
perf_dump_path_provided=batch.perf_dump_path is not None,
metrics=batch.metrics,
log_stage_start_end=not batch.is_warmup
and not (self.server_args and self.server_args.comfyui_mode),
perf_dump_path_provided=batch.perf_dump_path is not None,
):
result = self.forward(batch, server_args)

2
python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py
View File

@@ -232,7 +232,7 @@ class DecodingStage(PipelineStage):
trajectory_timesteps=batch.trajectory_timesteps,
trajectory_latents=batch.trajectory_latents,
trajectory_decoded=trajectory_decoded,
timings=batch.timings,
metrics=batch.metrics,
)
self.offload_model()

2
python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py
View File

@@ -69,7 +69,7 @@ class LTX2AVDecodingStage(DecodingStage):
trajectory_timesteps=batch.trajectory_timesteps,
trajectory_latents=batch.trajectory_latents,
trajectory_decoded=None,
timings=batch.timings,
metrics=batch.metrics,
)
# 2. Decode Audio

2
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
View File

@@ -1017,7 +1017,7 @@ class DenoisingStage(PipelineStage):
with StageProfiler(
f"denoising_step_{i}",
logger=logger,
timings=batch.timings,
metrics=batch.metrics,
perf_dump_path_provided=batch.perf_dump_path is not None,
):
t_int = int(t_host.item())

2
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py
View File

@@ -346,7 +346,7 @@ class LTX2AVDenoisingStage(DenoisingStage):
with StageProfiler(
f"denoising_step_{i}",
logger=logger,
timings=batch.timings,
metrics=batch.metrics,
perf_dump_path_provided=batch.perf_dump_path is not None,
):
t_int = int(t_host.item())

2
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py
View File

@@ -102,7 +102,7 @@ class DmdDenoisingStage(DenoisingStage):
with StageProfiler(
f"denoising_step_{i}",
logger=logger,
timings=batch.timings,
metrics=batch.metrics,
perf_dump_path_provided=batch.perf_dump_path is not None,
):
t_int = int(t.item())

6
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py
View File

@@ -399,7 +399,7 @@ class MOVADenoisingStage(PipelineStage):
getattr(batch, "extra_step_kwargs", None) or {},
)
timings = getattr(batch, "timings", None)
metrics = getattr(batch, "metrics", None)
perf_dump_path_provided = getattr(batch, "perf_dump_path", None) is not None
with self.progress_bar(total=total_steps) as progress_bar:
@@ -407,7 +407,7 @@ class MOVADenoisingStage(PipelineStage):
with StageProfiler(
f"denoising_step_{idx_step}",
logger=logger,
timings=timings,
metrics=metrics,
perf_dump_path_provided=perf_dump_path_provided,
):
pair_t = paired_timesteps[idx_step]
@@ -908,6 +908,6 @@ class MOVADecodingStage(PipelineStage):
output=video,
audio=audio,
audio_sample_rate=getattr(self.audio_vae, "sample_rate", None),
timings=batch.timings,
metrics=batch.metrics,
)
return output_batch

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

@@ -1,6 +1,7 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import dataclasses
import json
import logging
import os
import subprocess
import sys
@@ -15,7 +16,10 @@ from dateutil.tz import UTC
import sglang
import sglang.multimodal_gen.envs as envs
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import (
CYAN,
RESET,
_SGLDiffusionLogger,
get_is_main_process,
init_logger,
@@ -184,13 +188,13 @@ class StageProfiler:
self,
stage_name: str,
logger: _SGLDiffusionLogger,
timings: Optional["RequestMetrics"],
metrics: 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
self.metrics = metrics
self.logger = logger
self.start_time = 0.0
self.log_timing = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING
@@ -199,9 +203,12 @@ class StageProfiler:
def __enter__(self):
if self.log_stage_start_end:
self.logger.info(f"[{self.stage_name}] started...")
msg = f"[{self.stage_name}] started..."
if self.logger.isEnabledFor(logging.DEBUG):
msg += f" ({round(current_platform.get_available_gpu_memory(), 2)} GB left)"
self.logger.info(msg)
if (self.log_timing and self.timings) or self.log_stage_start_end:
if (self.log_timing and self.metrics) or self.log_stage_start_end:
if (
os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
and self.stage_name.startswith("denoising_step_")
@@ -213,7 +220,7 @@ class StageProfiler:
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if not ((self.log_timing and self.timings) or self.log_stage_start_end):
if not ((self.log_timing and self.metrics) or self.log_stage_start_end):
return False
if (
@@ -239,17 +246,17 @@ class StageProfiler:
f"[{self.stage_name}] finished in {execution_time_s:.4f} seconds",
)
if self.log_timing and self.timings:
if self.log_timing and self.metrics:
if "denoising_step_" in self.stage_name:
index = int(self.stage_name[len("denoising_step_") :])
self.timings.record_steps(index, execution_time_s)
self.metrics.record_steps(index, execution_time_s)
else:
self.timings.record_stage(self.stage_name, execution_time_s)
self.metrics.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(
self.metrics.record_memory_snapshot(
f"after_{self.stage_name}", snapshot
)
@@ -269,7 +276,7 @@ class PerformanceLogger:
def dump_benchmark_report(
cls,
file_path: str,
timings: "RequestMetrics",
metrics: "RequestMetrics",
meta: Optional[Dict[str, Any]] = None,
tag: str = "benchmark_dump",
):
@@ -279,25 +286,25 @@ class PerformanceLogger:
"""
formatted_steps = [
{"name": name, "duration_ms": duration_ms}
for name, duration_ms in timings.stages.items()
for name, duration_ms in metrics.stages.items()
]
denoise_steps_ms = [
{"step": idx, "duration_ms": duration_ms}
for idx, duration_ms in enumerate(timings.steps)
for idx, duration_ms in enumerate(metrics.steps)
]
memory_checkpoints = {
name: snapshot.to_dict()
for name, snapshot in timings.memory_snapshots.items()
for name, snapshot in metrics.memory_snapshots.items()
}
report = {
"timestamp": datetime.now(UTC).isoformat(),
"request_id": timings.request_id,
"request_id": metrics.request_id,
"commit_hash": get_git_commit_hash(),
"tag": tag,
"total_duration_ms": timings.total_duration_ms,
"total_duration_ms": metrics.total_duration_ms,
"steps": formatted_steps,
"denoise_steps_ms": denoise_steps_ms,
"memory_checkpoints": memory_checkpoints,
@@ -309,14 +316,14 @@ class PerformanceLogger:
os.makedirs(os.path.dirname(abs_path), exist_ok=True)
with open(abs_path, "w", encoding="utf-8") as f:
json.dump(report, f, indent=2)
logger.info(f"Metrics dumped to: {abs_path}")
logger.info(f"Metrics dumped to: {CYAN}{abs_path}{RESET}")
except IOError as e:
logger.error(f"Failed to dump metrics to {abs_path}: {e}")
@classmethod
def log_request_summary(
cls,
timings: "RequestMetrics",
metrics: "RequestMetrics",
tag: str = "total_inference_time",
):
"""logs the stage metrics and total duration for a completed request
@@ -326,21 +333,21 @@ class PerformanceLogger:
"""
formatted_stages = [
{"name": name, "execution_time_ms": duration_ms}
for name, duration_ms in timings.stages.items()
for name, duration_ms in metrics.stages.items()
]
memory_checkpoints = {
name: snapshot.to_dict()
for name, snapshot in timings.memory_snapshots.items()
for name, snapshot in metrics.memory_snapshots.items()
}
record = RequestPerfRecord(
timings.request_id,
metrics.request_id,
commit_hash=get_git_commit_hash(),
tag="pipeline_stage_metrics",
stages=formatted_stages,
steps=timings.steps,
total_duration_ms=timings.total_duration_ms,
steps=metrics.steps,
total_duration_ms=metrics.total_duration_ms,
memory_snapshots=memory_checkpoints,
)

2
python/sglang/multimodal_gen/test/server/perf_baselines.json
View File

@@ -542,7 +542,7 @@
"7": 93.97,
"8": 94.32
},
"expected_e2e_ms": 1192.92,
"expected_e2e_ms": 1292.92,
"expected_avg_denoise_ms": 83.75,
"expected_median_denoise_ms": 93.58
},