[diffusion] feat: layerwise NVTX markers for Nsight Systems profiling (#25683)

Co-authored-by: Mick <mickjagger19@icloud.com>

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.

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(

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"],

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(

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:

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,

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):

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",

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)

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
 

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()

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

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

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)