From 3142278c5f36ba7816c79392e1ce47d252b50f04 Mon Sep 17 00:00:00 2001 From: mispa-ms <81828223+mispa-ms@users.noreply.github.com> Date: Mon, 25 May 2026 20:24:27 -0700 Subject: [PATCH] [diffusion] feat: layerwise NVTX markers for Nsight Systems profiling (#25683) Co-authored-by: Mick --- .../memory_managers/component_manager.py | 105 ++++- .../executors/parallel_executor.py | 72 ++-- .../executors/pipeline_executor.py | 57 ++- .../pipelines_core/executors/sync_executor.py | 23 +- .../runtime/pipelines_core/stages/base.py | 57 ++- .../pipelines_core/stages/denoising.py | 94 +++-- .../pipelines_core/stages/text_encoding.py | 2 +- .../multimodal_gen/runtime/server_args.py | 14 + .../runtime/utils/nvtx_pytorch_hooks.py | 214 ++++++++++ .../test/unit/test_multi_output_grouping.py | 8 +- .../test/unit/test_nvtx_pytorch_hooks.py | 374 ++++++++++++++++++ .../unit/test_pipeline_stage_profiling.py | 4 +- .../srt/speculative/frozen_kv_mtp_worker.py | 10 +- .../unit/mem_cache/test_registry.py | 36 +- 14 files changed, 942 insertions(+), 128 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/utils/nvtx_pytorch_hooks.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py index 7501257d9..c5e25d073 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py @@ -25,6 +25,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_co ) 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.nvtx_pytorch_hooks import DiffusionNvtxHooks logger = init_logger(__name__) @@ -148,6 +149,10 @@ class ComponentResidencyManager: self._current_use_index: int = -1 self._active_use: ComponentUse | None = None self._active_use_module: nn.Module | None = None + self._active_nvtx_key: tuple[str, str, str | None] | None = None + self._nvtx_hooks_by_use_key: dict[ + tuple[str, str, str | None], tuple[int, DiffusionNvtxHooks] + ] = {} self._prefetched_use_keys: set[tuple[str, str, str | None]] = set() self._custom_strategies: dict[str, ComponentResidencyStrategy] = dict( pipeline.component_residency_strategies @@ -161,6 +166,7 @@ class ComponentResidencyManager: def refresh_pipeline(self, pipeline: ComponentResidencyPipeline) -> None: custom_strategies = dict(pipeline.component_residency_strategies) if pipeline is not self.pipeline: + self._remove_nvtx_hooks() self.strategy_for.cache_clear() self._should_keep_single_dit.cache_clear() self._active_use = None @@ -200,6 +206,7 @@ class ComponentResidencyManager: ) self._active_use = None self._active_use_module = None + self._disable_active_nvtx() self._current_use_index = -1 self._prefetched_use_keys.clear() self._uses_seen.clear() @@ -241,11 +248,11 @@ class ComponentResidencyManager: return self._trace("stage_exit", detail=f"index={stage_index}") - def before_use(self, use: ComponentUse) -> None: + def before_use(self, use: ComponentUse, module: nn.Module | None = None) -> None: """component use-site starts""" if not self.enabled: return - self.begin_use(use) + self.begin_use(use, module=module) def begin_use(self, use: ComponentUse, module: nn.Module | None = None) -> None: """Begin one sequential component use interval. this is idempotent @@ -255,8 +262,19 @@ class ComponentResidencyManager: 3. Wait until the current component is ready, then prefetch the next heavy use. """ if self._active_use is not None and self._same_use(self._active_use, use): + if self._use_key(self._active_use) != self._use_key(use): + self._mark_current_use(use) + self._active_use = use + self.state.current_use = use + self._enable_nvtx_for_use( + use, + module + or self._active_use_module + or self.get_module(use.component_name), + ) return if self._active_use is not None: + self._disable_active_nvtx() # finish previous active use self._finish_use( self._active_use, @@ -267,9 +285,10 @@ class ComponentResidencyManager: self._active_use_module = None self.state.current_use = None self._mark_current_use(use) - self._prepare_forward_use(use, module=module) + module = self._prepare_forward_use(use, module=module) self._active_use = use self._active_use_module = module + self._enable_nvtx_for_use(use, module) self._prefetch_next_memory_intensive_use() def end_use(self, use: ComponentUse, module: nn.Module | None = None) -> None: @@ -281,6 +300,7 @@ class ComponentResidencyManager: """ if self._active_use is None or not self._same_use(self._active_use, use): return + self._disable_active_nvtx() self._finish_use( self._active_use, module=self._active_use_module or module, @@ -323,6 +343,20 @@ class ComponentResidencyManager: return self._prepare_forward_use(use, module=module) + def remove_nvtx_hooks_for_module(self, module: nn.Module | None) -> None: + """Detach NVTX hooks before a component object is deleted or replaced.""" + if module is None: + return + module_id = id(module) + for key, (registered_id, hooks) in list(self._nvtx_hooks_by_use_key.items()): + if registered_id != module_id: + continue + if self._active_nvtx_key == key: + hooks.set_enabled(False) + self._active_nvtx_key = None + hooks.remove_hooks() + del self._nvtx_hooks_by_use_key[key] + def prefetch_checkpoint(self, anchor: ComponentUse | None = None) -> None: """Give the manager a timeline overlap point. @@ -341,6 +375,7 @@ class ComponentResidencyManager: if self._active_use is None: return active_use = self._active_use + self._disable_active_nvtx() self._finish_use( active_use, module=self._active_use_module, @@ -354,12 +389,12 @@ class ComponentResidencyManager: def _prepare_forward_use( self, use: ComponentUse, module: nn.Module | None = None - ) -> None: + ) -> nn.Module | None: """Prepare a component that is about to run and wait until it is ready.""" module = module or self.get_module(use.component_name) if module is None: self._trace("skip_missing", use) - return + return None strategy = self.strategy_for(use.component_name, module) self._uses_seen[use.component_name] = use self.state.current_use = use @@ -367,6 +402,66 @@ class ComponentResidencyManager: strategy.prepare_for_use(module, use, self.state) self._trace("wait", use, strategy, module) strategy.wait_for_use(module, use, self.state) + return module + + def _enable_nvtx_for_use( + self, use: ComponentUse, module: nn.Module | None = None + ) -> None: + if ( + not self.server_args.enable_layerwise_nvtx_marker + or self.state.batch_is_warmup + or not isinstance(module, nn.Module) + ): + self._disable_active_nvtx() + return + + key = self._use_key(use) + if self._active_nvtx_key != key: + self._disable_active_nvtx() + + module_id = id(module) + existing = self._nvtx_hooks_by_use_key.get(key) + if existing is None or existing[0] != module_id: + if existing is not None: + existing[1].remove_hooks() + self._nvtx_hooks_by_use_key.pop(key, None) + hooks = DiffusionNvtxHooks() + prefix = self._nvtx_prefix_for_use(use) + total = hooks.register_hooks(module, prefix=prefix) + if total == 0: + return + logger.debug( + "[component_residency] Registered NVTX hooks for %s on %d submodules", + prefix, + total, + ) + self._nvtx_hooks_by_use_key[key] = (module_id, hooks) + else: + hooks = existing[1] + + hooks.set_enabled(True) + self._active_nvtx_key = key + + def _disable_active_nvtx(self) -> None: + if self._active_nvtx_key is None: + return + existing = self._nvtx_hooks_by_use_key.get(self._active_nvtx_key) + if existing is not None: + existing[1].set_enabled(False) + self._active_nvtx_key = None + + def _remove_nvtx_hooks(self) -> None: + self._disable_active_nvtx() + for _, hooks in self._nvtx_hooks_by_use_key.values(): + hooks.remove_hooks() + self._nvtx_hooks_by_use_key.clear() + + @staticmethod + def _nvtx_prefix_for_use(use: ComponentUse) -> str: + parts = [use.stage_name, use.component_name] + if use.phase is not None and use.phase != use.component_name: + parts.append(use.phase) + return ".".join(parts) def _prefetch_use(self, use: ComponentUse) -> None: """Prepare a future component opportunistically without waiting. diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py index acaac0f68..48ffee232 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py @@ -4,7 +4,6 @@ from typing import Any, Callable, List import torch -from sglang.multimodal_gen.runtime.distributed import get_sp_group from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_cfg_group, get_classifier_free_guidance_rank, @@ -33,26 +32,6 @@ class ParallelExecutor(PipelineExecutor): """ - def collect_from_main(self, batches: list[Req]): - - # TODO: fix this condition - if self.server_args.sp_degree != 1: - sp_group = get_sp_group() - batches = broadcast_pyobj( - batches, - sp_group.rank, - sp_group.cpu_group, - src=sp_group.ranks[0], - ) - - if self.server_args.enable_cfg_parallel: - batches = broadcast_pyobj( - batches, - self.worker.cfg_group.rank, - self.worker.cfg_cpu_group, - src=self.worker.cfg_group.ranks[0], - ) - def _execute_stages( self, stages: List[PipelineStage], @@ -68,8 +47,9 @@ class ParallelExecutor(PipelineExecutor): cfg_group = get_cfg_group() group = get_world_group() - self.begin_component_residency_request(stages, batch, server_args) - try: + use_nvtx = self._should_use_stage_nvtx(batch, server_args) + + with self._component_residency_request(stages, batch, server_args): # TODO: decide when to gather on main when CFG_PARALLEL -> MAIN_RANK_ONLY for stage_index, stage in enumerate(stages): paradigm = stage.parallelism_type @@ -77,11 +57,14 @@ class ParallelExecutor(PipelineExecutor): if paradigm == StageParallelismType.MAIN_RANK_ONLY: if rank == 0: # Only main rank executes, others just wait - self.before_stage(stage, stage_index, batch, server_args) - batch = self.run_stage_with_context( - stage, batch, server_args, run_stage + batch = self._run_stage_with_executor_hooks( + stage, + stage_index, + batch, + server_args, + run_stage, + use_nvtx, ) - self.after_stage(stage_index) torch.distributed.barrier() elif paradigm == StageParallelismType.CFG_PARALLEL: @@ -95,28 +78,37 @@ class ParallelExecutor(PipelineExecutor): ) if rank != 0: batch = broadcasted_list[0] - self.before_stage(stage, stage_index, batch, server_args) - batch = self.run_stage_with_context( - stage, batch, server_args, run_stage + batch = self._run_stage_with_executor_hooks( + stage, + stage_index, + batch, + server_args, + run_stage, + use_nvtx, ) - self.after_stage(stage_index) torch.distributed.barrier() elif paradigm == StageParallelismType.REPLICATED: - self.before_stage(stage, stage_index, batch, server_args) - batch = self.run_stage_with_context( - stage, batch, server_args, run_stage + batch = self._run_stage_with_executor_hooks( + stage, + stage_index, + batch, + server_args, + run_stage, + use_nvtx, ) - self.after_stage(stage_index) elif paradigm == StageParallelismType.MAIN_RANK_ONLY_AND_SEND_TO_OTHERS: if rank == 0: # Only main rank executes, others just wait - self.before_stage(stage, stage_index, batch, server_args) - batch = self.run_stage_with_context( - stage, batch, server_args, run_stage + batch = self._run_stage_with_executor_hooks( + stage, + stage_index, + batch, + server_args, + run_stage, + use_nvtx, ) - self.after_stage(stage_index) torch.distributed.barrier() # Send batch to other ranks @@ -127,8 +119,6 @@ class ParallelExecutor(PipelineExecutor): if rank != 0: batch = broadcasted_list[0] torch.distributed.barrier() - finally: - self.finish_component_residency_request() return batch def execute( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py index e4960df97..be94ccde5 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py @@ -7,7 +7,7 @@ Base class for all pipeline executors. import contextlib from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING, Any, Callable, List import torch @@ -16,6 +16,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBa from sglang.multimodal_gen.runtime.platforms import current_platform 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.nvtx_pytorch_hooks import maybe_nvtx_range from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler @@ -53,7 +54,7 @@ class PipelineExecutor(ABC): def begin_component_residency_request( self, stages: List["PipelineStage"], - batch: Req, + batch: Any, server_args: ServerArgs, ) -> None: self.component_residency_manager.begin_request(stages, batch, server_args) @@ -62,7 +63,7 @@ class PipelineExecutor(ABC): self, stage: "PipelineStage", stage_index: int, - batch: Req, + batch: Any, server_args: ServerArgs, ) -> None: stage.set_component_residency_manager(self.component_residency_manager) @@ -76,6 +77,56 @@ class PipelineExecutor(ABC): def finish_component_residency_request(self) -> None: self.component_residency_manager.finish_request() + @contextlib.contextmanager + def _component_residency_request( + self, + stages: List["PipelineStage"], + payload: Any, + server_args: ServerArgs, + ): + self.begin_component_residency_request(stages, payload, server_args) + try: + yield + finally: + self.finish_component_residency_request() + + @staticmethod + def _is_warmup_payload(payload: Any) -> bool: + if isinstance(payload, list): + return bool(payload) and all( + getattr(item, "is_warmup", False) for item in payload + ) + return getattr(payload, "is_warmup", False) + + def _should_use_stage_nvtx(self, payload: Any, server_args: ServerArgs) -> bool: + return server_args.enable_layerwise_nvtx_marker and not self._is_warmup_payload( + payload + ) + + def _run_stage_with_executor_hooks( + self, + stage: "PipelineStage", + stage_index: int, + payload: Any, + server_args: ServerArgs, + run_stage: Callable[["PipelineStage", Any], Any], + use_nvtx: bool, + ) -> Any: + stage_name = stage._component_stage_name() + self.before_stage(stage, stage_index, payload, server_args) + with maybe_nvtx_range(f"stage_{stage_name}", use_nvtx): + payload = self.run_stage_with_context( + stage, payload, server_args, run_stage + ) + self.after_stage(stage_index) + return payload + + @staticmethod + def _step_stage_profiler() -> None: + profiler = SGLDiffusionProfiler.get_instance() + if profiler: + profiler.step_stage() + def execute_with_profiling( self, stages: List["PipelineStage"], diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py index d19036080..f9ece80da 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py @@ -9,7 +9,6 @@ from typing import Any, Callable, List from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import ( PipelineExecutor, - SGLDiffusionProfiler, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages import PipelineStage @@ -29,19 +28,19 @@ class SyncExecutor(PipelineExecutor): run_stage: Callable[[PipelineStage, Any], Any], ) -> Any: """Execute all pipeline stages sequentially and step the profiler.""" - self.begin_component_residency_request(stages, payload, server_args) - try: + + use_nvtx = self._should_use_stage_nvtx(payload, server_args) + with self._component_residency_request(stages, payload, server_args): for stage_index, stage in enumerate(stages): - self.before_stage(stage, stage_index, payload, server_args) - payload = self.run_stage_with_context( - stage, payload, server_args, run_stage + payload = self._run_stage_with_executor_hooks( + stage, + stage_index, + payload, + server_args, + run_stage, + use_nvtx, ) - self.after_stage(stage_index) - profiler = SGLDiffusionProfiler.get_instance() - if profiler: - profiler.step_stage() - finally: - self.finish_component_residency_request() + self._step_stage_profiler() return payload def run_profile_all_stages( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index 1044e95fc..c004ad728 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py @@ -59,6 +59,10 @@ class PipelineStage(StageDedupMixin, ABC): for a specific part of the process, such as prompt encoding, latent preparation, etc. """ + # Class-level default so subclasses that override __init__ without + # calling super().__init__() still see a consistent explicit-range gate. + _current_use_nvtx: bool = False + def __init__(self): self.server_args = get_global_server_args() self._component_residency_manager = None @@ -133,6 +137,12 @@ class PipelineStage(StageDedupMixin, ABC): ) def _active_component_stage_name(self) -> str: + """Stage name reported by the residency manager. + + Only valid between ``before_stage`` and ``after_stage``; outside + that window the manager state still holds the previous stage's + name. Use :meth:`_component_stage_name` for the static identity. + """ manager = getattr(self, "_component_residency_manager", None) manager_state = getattr(manager, "state", None) manager_stage_name = getattr(manager_state, "stage_name", None) @@ -206,6 +216,26 @@ class PipelineStage(StageDedupMixin, ABC): """Declares component uses of current stage for unified residency scheduling.""" return [] + def _apply_nvtx_gate(self, is_warmup: bool) -> bool: + """Resolve the per-request NVTX gate for explicit stage ranges. + + Layerwise module hooks are registered at component use-sites by + ``ComponentResidencyManager``. Stages use this value only for + explicit ``maybe_nvtx_range`` blocks. + """ + use_nvtx = self.server_args.enable_layerwise_nvtx_marker and not is_warmup + self._current_use_nvtx = use_nvtx + return use_nvtx + + @property + def current_use_nvtx(self) -> bool: + """Last resolved ``use_nvtx`` value from :meth:`_apply_nvtx_gate`. + + ``forward`` implementations can read this to gate explicit + ``maybe_nvtx_range`` blocks without re-evaluating the flag. + """ + return self._current_use_nvtx + # Default role affinity: ENCODER. Override in subclasses for DENOISING/DECODER. @property def role_affinity(self) -> RoleType: @@ -297,16 +327,23 @@ class PipelineStage(StageDedupMixin, ABC): logger.error("Input verification failed for %s: %s", stage_name, str(e)) raise - # Execute the actual stage logic with unified profiling - with StageProfiler( - stage_name, - logger=logger, - 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) + # Resolve the NVTX gate once per call. Component-level hooks are + # attached by the residency manager at the actual component use-site. + self._apply_nvtx_gate(batch.is_warmup) + + # Execute the actual stage logic with unified profiling. + try: + with StageProfiler( + stage_name, + logger=logger, + 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) + finally: + self._current_use_nvtx = False # Post-execution output verification try: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 761a09cfa..4b533ed90 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -98,6 +98,7 @@ from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin 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.nvtx_pytorch_hooks import maybe_nvtx_range from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler from sglang.multimodal_gen.utils import PRECISION_TO_TYPE, dict_to_3d_list @@ -970,8 +971,13 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): ) -> None: """Run one scheduler-backed denoising step in the shared base path. - Model-specific stages should override this instead of the whole loop whenever possible to achieve better performance + Model-specific stages should override this instead of the whole loop + whenever possible to achieve better performance. Overrides that bypass + ``_predict_noise_with_cfg`` / ``ctx.scheduler.step`` will lose the + inner ``predict_noise`` / ``scheduler_step`` NVTX markers emitted + below; mirror them in the override if those markers are needed. """ + use_nvtx = self.current_use_nvtx # 1. Prepare latent inputs in the model's compute dtype. latent_model_input = ctx.latents.to(ctx.target_dtype) if batch.image_latent is not None: @@ -998,32 +1004,34 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): ) # 4. Run the model prediction path, including CFG when enabled. - noise_pred = self._predict_noise_with_cfg( - current_model=step.current_model, - latent_model_input=latent_model_input, - timestep=timestep, - batch=batch, - timestep_index=step.step_index, - attn_metadata=step.attn_metadata, - target_dtype=ctx.target_dtype, - current_guidance_scale=step.current_guidance_scale, - cfg_policy=ctx.cfg_policy, - cfg_gate_state=ctx.extra.get("cfg_gate_state"), - server_args=server_args, - guidance=ctx.guidance, - latents=ctx.latents, - ) + with maybe_nvtx_range("predict_noise", use_nvtx): + noise_pred = self._predict_noise_with_cfg( + current_model=step.current_model, + latent_model_input=latent_model_input, + timestep=timestep, + batch=batch, + timestep_index=step.step_index, + attn_metadata=step.attn_metadata, + target_dtype=ctx.target_dtype, + current_guidance_scale=step.current_guidance_scale, + cfg_policy=ctx.cfg_policy, + cfg_gate_state=ctx.extra.get("cfg_gate_state"), + server_args=server_args, + guidance=ctx.guidance, + latents=ctx.latents, + ) if server_args.comfyui_mode: batch.noise_pred = noise_pred # 5. Advance the scheduler state with the predicted noise. - ctx.latents = ctx.scheduler.step( - model_output=noise_pred, - timestep=step.t_device, - sample=ctx.latents, - **ctx.extra_step_kwargs, - return_dict=False, - )[0] + with maybe_nvtx_range("scheduler_step", use_nvtx): + ctx.latents = ctx.scheduler.step( + model_output=noise_pred, + timestep=step.t_device, + sample=ctx.latents, + **ctx.extra_step_kwargs, + return_dict=False, + )[0] # 6. Re-apply any model-specific latent constraints after the update. ctx.latents = self.post_forward_for_ti2v_task( @@ -1144,6 +1152,10 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): "Memory before deallocating transformer: %s", torch.mps.current_allocated_memory(), ) + if self._component_residency_manager is not None: + self._component_residency_manager.remove_nvtx_hooks_for_module( + self.transformer + ) del self.transformer if pipeline is not None and "transformer" in pipeline.modules: del pipeline.modules["transformer"] @@ -1249,7 +1261,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): preferred_ready_after_request=component_name == "transformer", memory_intensive=True, ) - manager.begin_use(use) + manager.begin_use(use, module=current_model) def _select_and_manage_model( self, @@ -1334,19 +1346,33 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): # to avoid device-sync caused by timestep comparison timesteps_cpu = ctx.timesteps.cpu() num_timesteps = timesteps_cpu.shape[0] - with torch.autocast( - device_type=current_platform.device_type, - dtype=ctx.target_dtype, - enabled=ctx.autocast_enabled, + # Re-resolve the explicit-range gate so the per-step markers + # below honor this request's is_warmup state. Layer hooks are + # registered by the residency manager at the use-site. + use_nvtx = self._apply_nvtx_gate(ctx.is_warmup) + + with ( + torch.autocast( + device_type=current_platform.device_type, + dtype=ctx.target_dtype, + enabled=ctx.autocast_enabled, + ), + maybe_nvtx_range("denoising_loop", use_nvtx), ): with self.progress_bar(total=ctx.num_inference_steps) as progress_bar: for step_index, t_host in enumerate(timesteps_cpu): - with StageProfiler( - f"denoising_step_{step_index}", - logger=logger, - metrics=batch.metrics, - perf_dump_path_provided=batch.perf_dump_path is not None, - record_as_step=True, + # Use ``:.4g`` so flow-matching schedulers (e.g. FLUX) that + # use non-integer timesteps keep their precision in markers. + step_marker = f"denoising_step_{step_index}_t{t_host.item():.4g}" + with ( + maybe_nvtx_range(step_marker, use_nvtx), + StageProfiler( + f"denoising_step_{step_index}", + logger=logger, + metrics=batch.metrics, + perf_dump_path_provided=batch.perf_dump_path is not None, + record_as_step=True, + ), ): step = self._prepare_step_state( ctx, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py index c5e619561..c332c1652 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py @@ -429,7 +429,7 @@ class TextEncodingStage(PipelineStage): # TODO: Keep this begin-only interval until manager supports explicit # declared-use interval grouping. Wrapping each encoder call separately # can offload between positive and negative prompt encoding. - manager.before_use(use) + manager.begin_use(use, module=self.text_encoders[encoder_index]) def _forward_text_encoder(self, text_encoder, encoder_forward_kwargs): if not getattr(text_encoder, "uses_sglang_forward_context", True): diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 2ab9f66a5..8314bf369 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -219,6 +219,9 @@ class ServerArgs(DisaggArgsMixin): # Compilation enable_torch_compile: bool = False + # NVTX profiling + enable_layerwise_nvtx_marker: bool = False + # warmup warmup: bool = False warmup_resolutions: list[str] = None @@ -1190,6 +1193,17 @@ class ServerArgs(DisaggArgsMixin): + "However, will likely cause precision drifts. See (https://github.com/pytorch/pytorch/issues/145213)", ) + parser.add_argument( + "--enable-layerwise-nvtx-marker", + action=StoreBoolean, + default=ServerArgs.enable_layerwise_nvtx_marker, + help="Enable layerwise NVTX markers for profiling with Nsight Systems. " + "Adds NVTX ranges around each pipeline stage, the denoising loop, " + "every denoising step, the predict_noise / scheduler_step " + "sub-operations, and every transformer submodule forward (recursive). " + "Warmup steps are excluded to keep captured traces clean.", + ) + # warmup parser.add_argument( "--warmup", diff --git a/python/sglang/multimodal_gen/runtime/utils/nvtx_pytorch_hooks.py b/python/sglang/multimodal_gen/runtime/utils/nvtx_pytorch_hooks.py new file mode 100644 index 000000000..2b020fcd5 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/utils/nvtx_pytorch_hooks.py @@ -0,0 +1,214 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""PyTorch hooks for layerwise NVTX profiling in SGLang Diffusion. + +Mirrors the structure of ``sglang.srt.utils.nvtx_pytorch_hooks.PytHooks`` +but uses a compact ``{name} in={shapes}`` marker format that is well-suited +to DiT transformer blocks. See +``sglang.srt.utils.nvtx_pytorch_hooks`` for the LLM-runtime equivalent +that emits a richer per-layer parameter dict. +""" + +from __future__ import annotations + +import contextlib +from collections.abc import Iterator +from typing import Any + +import torch +import torch.cuda.nvtx as nvtx +from torch.utils.hooks import RemovableHandle + +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +# Module types that are too lightweight to warrant their own NVTX range. +# Skipping them keeps the captured timeline readable. +_DEFAULT_SKIP_TYPES: tuple[type, ...] = ( + torch.nn.Identity, + torch.nn.Dropout, + torch.nn.Dropout1d, + torch.nn.Dropout2d, + torch.nn.Dropout3d, +) + + +@contextlib.contextmanager +def maybe_nvtx_range(name: str, enabled: bool = True) -> Iterator[None]: + """Context manager that wraps a block of work in an NVTX range. + + Calls ``range_push`` / ``range_pop`` directly rather than going through + :func:`torch.cuda.nvtx.range`, which would otherwise interpret ``name`` as a + ``str.format`` template (so a literal ``{`` in the marker would raise + ``KeyError``). The ``range_pop`` is invoked from the ``finally`` clause, so + exceptions raised inside the ``with`` block cannot leak a half-open range. + + When ``enabled`` is ``False`` the function is a zero-cost no-op, suitable + for use under a per-request gate (e.g. warmup exclusion). + """ + if not enabled: + yield + return + nvtx.range_push(name) + try: + yield + finally: + nvtx.range_pop() + + +class DiffusionNvtxHooks: + """Register NVTX markers around each submodule forward pass. + + Each registered module emits an NVTX range covering its forward pass. + The range name encodes the qualified module name and the input tensor + shapes for downstream identification in Nsight Systems. + + Hook handles are retained so they can be removed via :meth:`remove_hooks`; + the same instance must not be reused across multiple model instances. + """ + + def __init__(self, skip_types: tuple[type, ...] = _DEFAULT_SKIP_TYPES) -> None: + self._skip_types = skip_types + self._module_to_name_map: dict[torch.nn.Module, str] = {} + self._hook_handles: list[RemovableHandle] = [] + # Caller must explicitly enable via ``set_enabled``. Default off + # so a forward that bypasses the component-use gate (e.g. an early + # warmup pass) cannot accidentally pollute the captured timeline. + self._enabled: bool = False + + def register_hooks( + self, + model: torch.nn.Module, + prefix: str = "", + ) -> int: + """Walk ``model`` and attach forward pre/post hooks to every module. + + Args: + model: Root module to instrument. + prefix: Optional name prefix prepended to every emitted range. + + Returns: + Number of modules instrumented. + + Notes: + Weight-tied or otherwise duplicated module instances are + skipped (the first occurrence wins) so each forward pass + produces exactly one NVTX range. + """ + instrumented = 0 + for name, module in model.named_modules(prefix=prefix): + if isinstance(module, self._skip_types): + continue + # Skip duplicate module instances (e.g., weight-tied layers). + # The check must happen before hook registration to avoid + # double-emitting NVTX ranges on the second occurrence. + if module in self._module_to_name_map: + logger.debug( + "NVTX: module %s already registered as '%s', skipping '%s'", + type(module).__name__, + self._module_to_name_map[module], + name, + ) + continue + self._module_to_name_map[module] = name + self._hook_handles.append( + module.register_forward_pre_hook( + self._forward_pre_hook, with_kwargs=True + ) + ) + # ``always_call=True`` (PyTorch 2.0+) guarantees the post-hook + # still fires when ``forward`` raises, so an OOM or assertion + # inside the wrapped module cannot leak a half-open NVTX range. + self._hook_handles.append( + module.register_forward_hook(self._forward_hook, always_call=True) + ) + instrumented += 1 + return instrumented + + def remove_hooks(self) -> None: + """Remove every hook registered by this instance. + + Safe to call multiple times; subsequent calls are no-ops. The + bookkeeping is cleared in a ``finally`` so a misbehaving + ``handle.remove()`` cannot leave the instance with stale + handles or name-map entries. + """ + try: + for handle in self._hook_handles: + handle.remove() + finally: + self._hook_handles.clear() + self._module_to_name_map.clear() + + def set_enabled(self, enabled: bool) -> None: + """Toggle whether the registered hooks emit NVTX ranges. + + When disabled, both the pre- and post-hooks early-return, so each + forward produces a matched (push, pop) pair of "no-ops" — no range + leak and no half-open range across the toggle. + """ + self._enabled = enabled + + # ------------------------------------------------------------------ hooks + + def _forward_pre_hook( + self, + module: torch.nn.Module, + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> None: + if not self._enabled: + return + name = self._module_to_name_map.get(module, "unknown") + shapes = _collect_input_shapes(args, kwargs) + marker = f"{name} in={shapes}" if shapes else name + nvtx.range_push(marker) + + def _forward_hook( + self, + module: torch.nn.Module, + _args: Any, + _output: Any, + ) -> None: + if not self._enabled: + return + nvtx.range_pop() + + +def _collect_input_shapes( + args: tuple[Any, ...], kwargs: dict[str, Any] | None = None +) -> list[list[int]]: + """Best-effort extraction of input tensor shapes for marker labels. + + Walks positional ``args`` and keyword ``kwargs`` values, recursing into + lists and tuples (so DiT inputs like ``image_rotary_emb=(cos, sin)`` are + captured). Non-tensor scalars, ``None``, dicts, and arbitrary objects are + silently skipped. + """ + shapes: list[list[int]] = [] + _append_tensor_shapes(args, shapes) + if kwargs: + _append_tensor_shapes(tuple(kwargs.values()), shapes) + return shapes + + +def _append_tensor_shapes(items: Any, shapes: list[list[int]]) -> None: + if isinstance(items, torch.Tensor): + shapes.append(list(items.size())) + return + if isinstance(items, (list, tuple)): + for item in items: + _append_tensor_shapes(item, shapes) diff --git a/python/sglang/multimodal_gen/test/unit/test_multi_output_grouping.py b/python/sglang/multimodal_gen/test/unit/test_multi_output_grouping.py index 34aaccb03..3f0edafb3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_multi_output_grouping.py +++ b/python/sglang/multimodal_gen/test/unit/test_multi_output_grouping.py @@ -22,7 +22,9 @@ class CountingDedupStage(PipelineStage): deduplicated_extra_tensor_tree_output_keys = ("mu",) def __init__(self): - self.server_args = SimpleNamespace(comfyui_mode=True) + self.server_args = SimpleNamespace( + comfyui_mode=True, enable_layerwise_nvtx_marker=False + ) self.forward_calls = 0 def build_dedup_fingerprint(self, batch: Req, server_args): @@ -40,7 +42,9 @@ class CountingDedupStage(PipelineStage): class CountingLatentStage(LatentPreparationStage): def __init__(self): - self.server_args = SimpleNamespace(comfyui_mode=True) + self.server_args = SimpleNamespace( + comfyui_mode=True, enable_layerwise_nvtx_marker=False + ) self.prepare_group_calls = 0 self.forward_calls = 0 diff --git a/python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py b/python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py new file mode 100644 index 000000000..5483ceef5 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py @@ -0,0 +1,374 @@ +"""Unit tests for ``sglang.multimodal_gen.runtime.utils.nvtx_pytorch_hooks``. + +These tests cover the CPU-only surface: the ``maybe_nvtx_range`` helper, +``DiffusionNvtxHooks.register_hooks`` / ``remove_hooks`` lifecycle, and the +shape-collection helper. The actual ``nvtx.range_push`` / ``range_pop`` calls +require CUDA and are exercised end-to-end by Nsight-Systems profiling runs. +""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( + ComponentResidencyManager, + ComponentUse, +) +from sglang.multimodal_gen.runtime.utils import nvtx_pytorch_hooks +from sglang.multimodal_gen.runtime.utils.nvtx_pytorch_hooks import ( + DiffusionNvtxHooks, + _collect_input_shapes, + maybe_nvtx_range, +) + + +class TestMaybeNvtxRange(unittest.TestCase): + def test_disabled_returns_noop_context_manager(self) -> None: + ran = False + with maybe_nvtx_range("never", enabled=False): + ran = True + self.assertTrue(ran) + + def test_disabled_propagates_exception(self) -> None: + with self.assertRaises(RuntimeError): + with maybe_nvtx_range("never", enabled=False): + raise RuntimeError("boom") + + def test_disabled_does_not_call_nvtx(self) -> None: + with ( + patch.object(nvtx_pytorch_hooks.nvtx, "range_push") as push, + patch.object(nvtx_pytorch_hooks.nvtx, "range_pop") as pop, + ): + with maybe_nvtx_range("never", enabled=False): + pass + push.assert_not_called() + pop.assert_not_called() + + def test_enabled_calls_matched_push_pop(self) -> None: + with ( + patch.object(nvtx_pytorch_hooks.nvtx, "range_push") as push, + patch.object(nvtx_pytorch_hooks.nvtx, "range_pop") as pop, + ): + with maybe_nvtx_range("stage_X", enabled=True): + pass + push.assert_called_once_with("stage_X") + pop.assert_called_once_with() + + def test_enabled_pops_on_exception(self) -> None: + with ( + patch.object(nvtx_pytorch_hooks.nvtx, "range_push") as push, + patch.object(nvtx_pytorch_hooks.nvtx, "range_pop") as pop, + ): + with self.assertRaises(RuntimeError): + with maybe_nvtx_range("stage_X", enabled=True): + raise RuntimeError("boom") + push.assert_called_once_with("stage_X") + pop.assert_called_once_with() + + def test_marker_with_braces_does_not_raise(self) -> None: + """Regression: torch.cuda.nvtx.range() str-formats its argument, + which would raise on a marker containing a literal ``{``. The helper + calls range_push directly to sidestep that.""" + with ( + patch.object(nvtx_pytorch_hooks.nvtx, "range_push"), + patch.object(nvtx_pytorch_hooks.nvtx, "range_pop"), + ): + with maybe_nvtx_range("layer in={1, 2, 3}", enabled=True): + pass + + +class _TinyBlock(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(4, 4) + self.norm = torch.nn.LayerNorm(4) + # Dropout is in _DEFAULT_SKIP_TYPES and must not be instrumented. + self.drop = torch.nn.Dropout(p=0.0) + + +class TestDiffusionNvtxHooks(unittest.TestCase): + def test_register_hooks_counts_non_skipped_submodules(self) -> None: + block = _TinyBlock() + hooks = DiffusionNvtxHooks() + # 4 modules total (block, linear, norm, drop); drop is skipped. + self.assertEqual(hooks.register_hooks(block, prefix="block"), 3) + # 2 hooks (pre + post) registered per instrumented module. + self.assertEqual(len(hooks._hook_handles), 6) + + def test_register_hooks_skips_duplicate_instances(self) -> None: + shared = torch.nn.Linear(4, 4) + + class TiedModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.a = shared + self.b = shared + + model = TiedModel() + hooks = DiffusionNvtxHooks() + # Root + 1 unique linear (the second occurrence is skipped). + self.assertEqual(hooks.register_hooks(model), 2) + + def test_remove_hooks_is_idempotent(self) -> None: + block = _TinyBlock() + hooks = DiffusionNvtxHooks() + hooks.register_hooks(block) + hooks.remove_hooks() + self.assertEqual(hooks._hook_handles, []) + self.assertEqual(hooks._module_to_name_map, {}) + # Second call is a no-op, not an error. + hooks.remove_hooks() + + def test_set_enabled_false_suppresses_nvtx_calls(self) -> None: + """When disabled, neither pre- nor post-hook should call nvtx — + guarantees no half-open push without a matching pop.""" + hooks = DiffusionNvtxHooks() + dummy = torch.nn.Linear(2, 2) + hooks._module_to_name_map[dummy] = "dummy" + hooks.set_enabled(False) + with ( + patch.object(nvtx_pytorch_hooks.nvtx, "range_push") as push, + patch.object(nvtx_pytorch_hooks.nvtx, "range_pop") as pop, + ): + hooks._forward_pre_hook(dummy, (torch.zeros(2),), {}) + hooks._forward_hook(dummy, (), None) + push.assert_not_called() + pop.assert_not_called() + + def test_set_enabled_true_emits_matched_push_pop(self) -> None: + """When enabled, a forward pre/post pair emits exactly one push + and one pop with the qualified module name as the marker.""" + hooks = DiffusionNvtxHooks() + dummy = torch.nn.Linear(2, 2) + hooks._module_to_name_map[dummy] = "dummy" + hooks.set_enabled(True) + with ( + patch.object(nvtx_pytorch_hooks.nvtx, "range_push") as push, + patch.object(nvtx_pytorch_hooks.nvtx, "range_pop") as pop, + ): + hooks._forward_pre_hook(dummy, (torch.zeros(2, 3),), {}) + hooks._forward_hook(dummy, (), None) + push.assert_called_once() + marker = push.call_args.args[0] + self.assertIn("dummy", marker) + self.assertIn("[2, 3]", marker) + pop.assert_called_once_with() + + def test_default_enabled_is_false(self) -> None: + """Default off so an unguarded forward (e.g. early warmup) cannot + emit ranges; the caller must explicitly enable via set_enabled.""" + self.assertFalse(DiffusionNvtxHooks()._enabled) + + def test_post_hook_fires_on_forward_exception(self) -> None: + """Regression: ``always_call=True`` on the registered post-hook + guarantees ``range_pop`` runs even when the wrapped ``forward`` + raises. Without it an OOM (or any other forward-time exception) + would leak a half-open NVTX range.""" + + class _RaisingModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + raise RuntimeError("simulated forward exception") + + model = _RaisingModule() + hooks = DiffusionNvtxHooks() + hooks.register_hooks(model, prefix="raising") + hooks.set_enabled(True) + with ( + patch.object(nvtx_pytorch_hooks.nvtx, "range_push") as push, + patch.object(nvtx_pytorch_hooks.nvtx, "range_pop") as pop, + ): + with self.assertRaises(RuntimeError): + model(torch.zeros(2)) + # One push from the pre-hook, one pop from the post-hook fired via + # always_call=True; they must match to keep the stack balanced. + self.assertEqual(push.call_count, pop.call_count) + self.assertEqual(push.call_count, 1) + + +class _NoOpResidencyStrategy: + name = "noop" + + def prepare_for_use(self, module, use, state) -> None: + pass + + def wait_for_use(self, module, use, state) -> None: + pass + + def finish_use(self, module, use, state) -> None: + pass + + def finish_request(self, module, use, state, *, preferred: bool) -> None: + pass + + def prefetch_for_use(self, module, use, state) -> bool: + return False + + def prepare_after_request(self, module, use, state) -> None: + pass + + +def _test_manager( + modules: dict[str, torch.nn.Module], + *, + enable_flag: bool = True, + is_warmup: bool = False, +) -> ComponentResidencyManager: + pipeline = SimpleNamespace( + modules=modules, + _stage_name_mapping={}, + component_residency_strategies={}, + ) + server_args = SimpleNamespace(enable_layerwise_nvtx_marker=enable_flag) + manager = ComponentResidencyManager(pipeline, server_args) + manager.state.batch_is_warmup = is_warmup + manager.strategy_for = lambda _component_name, _module: _NoOpResidencyStrategy() + return manager + + +class TestComponentResidencyNvtxHooks(unittest.TestCase): + def test_disabled_flag_is_noop(self) -> None: + module = torch.nn.Linear(2, 2) + manager = _test_manager({"linear": module}, enable_flag=False) + manager.begin_use(ComponentUse("Stage", "linear"), module=module) + self.assertEqual(manager._nvtx_hooks_by_use_key, {}) + + def test_warmup_is_noop(self) -> None: + module = torch.nn.Linear(2, 2) + manager = _test_manager({"linear": module}, is_warmup=True) + manager.begin_use(ComponentUse("Stage", "linear"), module=module) + self.assertEqual(manager._nvtx_hooks_by_use_key, {}) + + def test_begin_use_registers_and_enables_component_hooks(self) -> None: + module = torch.nn.Linear(2, 2) + manager = _test_manager({"linear": module}) + use = ComponentUse("Stage", "linear") + + manager.begin_use(use, module=module) + + _, hooks = manager._nvtx_hooks_by_use_key[("Stage", "linear", None)] + self.assertTrue(hooks._enabled) + self.assertIn(module, hooks._module_to_name_map) + self.assertTrue(hooks._module_to_name_map[module].startswith("Stage.linear")) + + def test_end_use_disables_component_hooks(self) -> None: + module = torch.nn.Linear(2, 2) + manager = _test_manager({"linear": module}) + use = ComponentUse("Stage", "linear") + + manager.begin_use(use, module=module) + _, hooks = manager._nvtx_hooks_by_use_key[("Stage", "linear", None)] + manager.end_use(use, module=module) + + self.assertFalse(hooks._enabled) + self.assertIsNone(manager._active_nvtx_key) + + def test_remove_nvtx_hooks_for_module_drops_stale_reference(self) -> None: + module = torch.nn.Linear(2, 2) + manager = _test_manager({"linear": module}) + use = ComponentUse("Stage", "linear") + + manager.begin_use(use, module=module) + _, hooks = manager._nvtx_hooks_by_use_key[("Stage", "linear", None)] + manager.remove_nvtx_hooks_for_module(module) + + self.assertEqual(manager._nvtx_hooks_by_use_key, {}) + self.assertEqual(hooks._module_to_name_map, {}) + self.assertIsNone(manager._active_nvtx_key) + + def test_re_registers_when_module_identity_changes(self) -> None: + use = ComponentUse("Stage", "linear") + first_module = torch.nn.Linear(2, 2) + manager = _test_manager({"linear": first_module}) + + manager.begin_use(use, module=first_module) + _, first_hooks = manager._nvtx_hooks_by_use_key[("Stage", "linear", None)] + manager.end_use(use, module=first_module) + + second_module = torch.nn.Linear(2, 2) + manager.pipeline.modules["linear"] = second_module + manager.begin_use(use, module=second_module) + + _, second_hooks = manager._nvtx_hooks_by_use_key[("Stage", "linear", None)] + self.assertIsNot(second_hooks, first_hooks) + self.assertEqual(first_hooks._module_to_name_map, {}) + self.assertIn(second_module, second_hooks._module_to_name_map) + + def test_same_component_in_different_stages_switches_active_prefix(self) -> None: + shared = torch.nn.Linear(2, 2) + manager = _test_manager({"vae": shared}) + first_use = ComponentUse("ImageVAEEncodingStage", "vae") + second_use = ComponentUse("DecodingStage", "vae") + + manager.begin_use(first_use, module=shared) + _, first_hooks = manager._nvtx_hooks_by_use_key[ + ("ImageVAEEncodingStage", "vae", None) + ] + manager.begin_use(second_use, module=shared) + _, second_hooks = manager._nvtx_hooks_by_use_key[("DecodingStage", "vae", None)] + + self.assertFalse(first_hooks._enabled) + self.assertTrue(second_hooks._enabled) + self.assertTrue( + second_hooks._module_to_name_map[shared].startswith("DecodingStage.vae") + ) + + def test_pipeline_stage_call_sets_explicit_range_gate_before_forward(self) -> None: + from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( + PipelineStage, + ) + + class _Spy(PipelineStage): + def __init__(self) -> None: + self.server_args = type( + "Args", + (), + { + "enable_layerwise_nvtx_marker": True, + "comfyui_mode": False, + }, + )() + self._component_residency_manager = None + self._registered_stage_name = None + self._profile_stage_name = None + self._current_use_nvtx = False + self.use_nvtx_during_forward: bool | None = None + + def forward(self, batch, server_args): + self.use_nvtx_during_forward = self.current_use_nvtx + return batch + + class _Batch: + is_warmup = False + metrics = None + perf_dump_path = None + + spy = _Spy() + spy(_Batch(), spy.server_args) + self.assertTrue(spy.use_nvtx_during_forward) + self.assertFalse(spy.current_use_nvtx) + + +class TestCollectInputShapes(unittest.TestCase): + def test_flat_positional_tensors(self) -> None: + a = torch.zeros(2, 3) + b = torch.zeros(4) + self.assertEqual(_collect_input_shapes((a, b)), [[2, 3], [4]]) + + def test_kwarg_tensors_are_captured(self) -> None: + kw = {"hidden_states": torch.zeros(1, 4)} + self.assertEqual(_collect_input_shapes((), kw), [[1, 4]]) + + def test_nested_tuple_kwarg_recurses(self) -> None: + rope = (torch.zeros(8, 16), torch.zeros(8, 16)) + kw = {"image_rotary_emb": rope} + self.assertEqual(_collect_input_shapes((), kw), [[8, 16], [8, 16]]) + + def test_non_tensor_values_are_skipped(self) -> None: + kw = {"scale": 1.0, "use_cache": True, "extras": None} + self.assertEqual(_collect_input_shapes((42, "s"), kw), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_pipeline_stage_profiling.py b/python/sglang/multimodal_gen/test/unit/test_pipeline_stage_profiling.py index edb5e141b..a461aac32 100644 --- a/python/sglang/multimodal_gen/test/unit/test_pipeline_stage_profiling.py +++ b/python/sglang/multimodal_gen/test/unit/test_pipeline_stage_profiling.py @@ -7,7 +7,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineSta class NamedNoOpStage(PipelineStage): def __init__(self): - self.server_args = SimpleNamespace(comfyui_mode=True) + self.server_args = SimpleNamespace( + comfyui_mode=True, enable_layerwise_nvtx_marker=False + ) def forward(self, batch: Req, server_args) -> Req: return batch diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py index 4bad85187..d1e8daf93 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py @@ -398,8 +398,9 @@ class FrozenKVMTPWorker(TpModelWorker): forward_batch.mm_input_embeds = mm_input_embeds self._set_positions(forward_batch) self._init_frozen_kv_metadata(forward_batch) - with self._target_kv_pool_view(forward_batch), forward_context( - ForwardContext(attn_backend=self.draft_attn_backend) + with ( + self._target_kv_pool_view(forward_batch), + forward_context(ForwardContext(attn_backend=self.draft_attn_backend)), ): logits_output = self.draft_model_runner.forward( forward_batch, skip_attn_backend_init=True @@ -682,8 +683,9 @@ class FrozenKVMTPWorker(TpModelWorker): forward_batch.spec_info.hidden_states = hidden_states self._set_positions(forward_batch) - with self._target_kv_pool_view(forward_batch), forward_context( - ForwardContext(attn_backend=self.draft_attn_backend) + with ( + self._target_kv_pool_view(forward_batch), + forward_context(ForwardContext(attn_backend=self.draft_attn_backend)), ): logits_output = self.draft_model_runner.forward( forward_batch, skip_attn_backend_init=True diff --git a/test/registered/unit/mem_cache/test_registry.py b/test/registered/unit/mem_cache/test_registry.py index 302d61f11..026ee98c4 100644 --- a/test/registered/unit/mem_cache/test_registry.py +++ b/test/registered/unit/mem_cache/test_registry.py @@ -178,12 +178,15 @@ class TestDefaultRadixCacheFactory(CustomTestCase): # we inject a stand-in module rather than letting patch() trigger # the real import. fake_module = MagicMock() - with patch( - "sglang.srt.mem_cache.registry.envs.SGLANG_EXPERIMENTAL_CPP_RADIX_TREE.get", - return_value=True, - ), patch.dict( - "sys.modules", - {"sglang.srt.mem_cache.radix_cache_cpp": fake_module}, + with ( + patch( + "sglang.srt.mem_cache.registry.envs.SGLANG_EXPERIMENTAL_CPP_RADIX_TREE.get", + return_value=True, + ), + patch.dict( + "sys.modules", + {"sglang.srt.mem_cache.radix_cache_cpp": fake_module}, + ), ): result = default_radix_cache_factory(ctx) fake_module.RadixCacheCpp.assert_called_once_with( @@ -196,15 +199,18 @@ class TestDefaultRadixCacheFactory(CustomTestCase): # Shim both factory imports — each transitively loads sgl_kernel. fake_components = MagicMock() fake_radix = MagicMock() - with patch( - "sglang.srt.mem_cache.registry.envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get", - return_value=True, - ), patch.dict( - "sys.modules", - { - "sglang.srt.mem_cache.unified_cache_components": fake_components, - "sglang.srt.mem_cache.unified_radix_cache": fake_radix, - }, + with ( + patch( + "sglang.srt.mem_cache.registry.envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get", + return_value=True, + ), + patch.dict( + "sys.modules", + { + "sglang.srt.mem_cache.unified_cache_components": fake_components, + "sglang.srt.mem_cache.unified_radix_cache": fake_radix, + }, + ), ): result = default_radix_cache_factory(ctx) fake_radix.UnifiedRadixCache.assert_called_once_with(ctx.params)