[hardware, profile] add optional torch.profiler activity hook to Plat…

tests/plugin/test_platform_abstraction.py

@@ -208,5 +208,14 @@ def test_unregistered_platform_raises(self):
             _create_platform("nonexistent_platform")
 
 
+class TestTorchProfilerHooks:
+    """Test the default (no-op) behavior of the torch.profiler plugin hooks."""
+
+    def test_defaults_are_none(self):
+        platform = _make_mock_platform("mock_torch_profiler")()
+        assert platform.torch_profiler_activity() is None
+        assert platform.torch_profiler_content_name() is None
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

tests/utils/test_torch_profile.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import os
 import tempfile
 import unittest
+from unittest import mock
 from unittest.mock import MagicMock, patch
 
 import torch
@@ -646,6 +648,119 @@ def test_contents_are_not_rewritten(self, mock_profile):
         self.assertEqual(contents, ["cuda"])
 
 
+class TestPluginProfilerActivityHook(unittest.TestCase):
+    """A platform-supplied ``torch.profiler`` activity is added only when both
+    ``torch_profiler_activity()`` and ``torch_profiler_content_name()`` return non-``None``, and
+    (if ``contents`` is non-empty) the content name is one of the requested keywords. CUDA must
+    stay enabled when the user explicitly asks for it, even on a platform that also exposes its
+    own activity hook (regression guard for the plugin-device-availability vs.
+    plugin-device-requested mixup)."""
+
+    # A vendor-neutral stand-in: the code under test never interprets this value, it only
+    # stores/compares it, so any sentinel distinct from the real ProfilerActivity members works.
+    _MOCK_ACTIVITY = object()
+
+    def _mock_platform(self, activity, content_name):
+        platform = MagicMock()
+        platform.torch_profiler_activity.return_value = activity
+        platform.torch_profiler_content_name.return_value = content_name
+        return platform
+
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_added_for_empty_or_matching_contents(self, mock_profile, mock_get_platform):
+        mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, "mock_device")
+        for contents in ([], ["mock_device"]):
+            with self.subTest(contents=contents):
+                mock_profile.reset_mock()
+                get_torch_profiler(contents=list(contents), save_path="/tmp/test", rank=0)
+                activities = mock_profile.call_args[1]["activities"]
+                self.assertIn(self._MOCK_ACTIVITY, activities)
+
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_excluded_for_non_matching_contents(self, mock_profile, mock_get_platform):
+        mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, "mock_device")
+        get_torch_profiler(contents=["stack"], save_path="/tmp/test", rank=0)
+        activities = mock_profile.call_args[1]["activities"]
+        self.assertNotIn(self._MOCK_ACTIVITY, activities)
+
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_excluded_when_content_name_missing(self, mock_profile, mock_get_platform):
+        # Regression guard: an activity hook without a paired content name must never be added,
+        # even for empty `contents`.
+        mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, None)
+        get_torch_profiler(contents=[], save_path="/tmp/test", rank=0)
+        activities = mock_profile.call_args[1]["activities"]
+        self.assertNotIn(self._MOCK_ACTIVITY, activities)
+
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_excluded_when_activity_missing(self, mock_profile, mock_get_platform):
+        mock_get_platform.return_value = self._mock_platform(None, "mock_device")
+        get_torch_profiler(contents=["mock_device"], save_path="/tmp/test", rank=0)
+        activities = mock_profile.call_args[1]["activities"]
+        self.assertEqual(activities, [torch.profiler.ProfilerActivity.CPU])
+
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_explicit_cuda_wins_over_available_plugin_device(self, mock_profile, mock_get_platform):
+        # A plugin platform merely exposing an activity must not steal `contents=["cuda"]` away
+        # from CUDA: the plugin device has to be *requested*, not just *available*.
+        mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, "mock_device")
+        get_torch_profiler(contents=["cuda"], save_path="/tmp/test", rank=0)
+        activities = mock_profile.call_args[1]["activities"]
+        self.assertIn(torch.profiler.ProfilerActivity.CUDA, activities)
+        self.assertNotIn(self._MOCK_ACTIVITY, activities)
+
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_warns_when_cuda_is_requested_but_superseded(self, mock_profile, mock_get_platform):
+        # contents=["cuda", <plugin device>] is self-contradictory: only one activity is ever
+        # recorded, and there is no active CUDA platform here, so "cuda" is silently dead unless
+        # we say so.
+        mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, "mock_device")
+        with self.assertLogs("verl.utils.profiler.torch_profile", level="WARNING") as log_ctx:
+            get_torch_profiler(contents=["cuda", "mock_device"], save_path="/tmp/test", rank=0)
+        self.assertTrue(any("ignored" in record.getMessage() for record in log_ctx.records))
+
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_no_warning_without_explicit_cuda(self, mock_profile, mock_get_platform):
+        mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, "mock_device")
+        logger = logging.getLogger("verl.utils.profiler.torch_profile")
+        with mock.patch.object(logger, "warning") as mock_warning:
+            get_torch_profiler(contents=["mock_device"], save_path="/tmp/test", rank=0)
+        mock_warning.assert_not_called()
+
+
+class TestPluginContentNameConfigValidation(unittest.TestCase):
+    """``TorchProfilerToolConfig.__post_init__`` must only whitelist a platform's
+    ``torch_profiler_content_name()`` keyword when ``torch_profiler_activity()`` is also
+    non-``None`` -- otherwise a config naming that keyword passes validation here but
+    ``get_torch_profiler()`` silently drops it (see ``TestPluginProfilerActivityHook``)."""
+
+    _MOCK_ACTIVITY = object()
+
+    def _mock_platform(self, activity, content_name):
+        platform = MagicMock()
+        platform.torch_profiler_activity.return_value = activity
+        platform.torch_profiler_content_name.return_value = content_name
+        return platform
+
+    @patch("verl.utils.profiler.config.get_platform")
+    def test_content_name_allowed_when_activity_present(self, mock_get_platform):
+        mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, "mock_device")
+        TorchProfilerToolConfig(contents=["mock_device"], discrete=False)
+
+    @patch("verl.utils.profiler.config.get_platform")
+    def test_content_name_rejected_when_activity_missing(self, mock_get_platform):
+        mock_get_platform.return_value = self._mock_platform(None, "mock_device")
+        with self.assertRaises(AssertionError):
+            TorchProfilerToolConfig(contents=["mock_device"], discrete=False)
+
+
 def _role_profiler_omegaconf(tool="torch", enable=True, discrete=False, contents=("cpu", "cuda")):
     """Mimic a per-role ``profiler`` OmegaConf sub-tree (identical across ref/ref.yaml and
     critic/critic.yaml).

verl/plugin/platform/README.md

@@ -230,6 +230,8 @@ VERL_USE_EXTERNAL_PLUGINS=pkg1,pkg2  # only load named entry_points
 | **Profiling** | `nvtx_range(msg)` | Context manager for profiler ranges |
 | | `profiler_start()` | Start device profiler |
 | | `profiler_stop()` | Stop device profiler |
+| | `torch_profiler_activity()` | This platform's `torch.profiler.ProfilerActivity`, or `None` to opt out |
+| | `torch_profiler_content_name()` | `contents` keyword that requests it (e.g. `'xpu'`); must be paired with `torch_profiler_activity()` |
 | **Low-level** | `cudart()` | CUDA runtime API object or `None` |
 | | `check_smi_command(cmd)` | Run SMI command and check exit code |
 

verl/plugin/platform/platform_base.py

@@ -14,7 +14,10 @@
 import subprocess
 from contextlib import contextmanager
 from types import ModuleType
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    import torch
 
 
 class PlatformBase(abc.ABC):
@@ -189,6 +192,27 @@ def profiler_stop(self) -> None:
         """Stop the device profiler (no-op on unsupported platforms)."""
         ...
 
+    def torch_profiler_activity(self) -> Optional["torch.profiler.ProfilerActivity"]:
+        """Return this platform's ``torch.profiler`` device activity, if any.
+
+        Lets a platform request its own device-activity be collected by
+        ``verl/utils/profiler/torch_profile.py::get_torch_profiler`` when the
+        matching name (see :meth:`torch_profiler_content_name`) appears in
+        ``profiler.tool_config.torch.contents``. Return ``None`` (default) if
+        this platform has no ``torch.profiler.ProfilerActivity`` member (the
+        ``cuda`` activity is handled separately and is unaffected by this hook).
+        """
+        return None
+
+    def torch_profiler_content_name(self) -> Optional[str]:
+        """Return the ``contents`` keyword that requests this platform's activity.
+
+        E.g. ``"xpu"`` for Intel XPU. Paired with :meth:`torch_profiler_activity`;
+        both must be non-``None`` for the activity to be added. Return ``None``
+        (default) to opt out.
+        """
+        return None
+
     # ------------------------------------------------------------------
     # vllm integration
     # ------------------------------------------------------------------

verl/utils/profiler/config.py

@@ -24,6 +24,7 @@
 from omegaconf import MISSING
 
 from verl.base_config import BaseConfig
+from verl.plugin.platform import get_platform
 
 
 @dataclass
@@ -127,6 +128,10 @@ class TorchProfilerToolConfig(BaseConfig):
     def __post_init__(self) -> None:
         """config validation logics go here"""
         __support_contents = ["cuda", "cpu", "memory", "shapes", "stack"]
+        plugin_content = get_platform().torch_profiler_content_name()
+        plugin_activity = get_platform().torch_profiler_activity()
+        if plugin_content is not None and plugin_activity is not None:
+            __support_contents = [*__support_contents, plugin_content]
         for content in self.contents:
             assert content in __support_contents, (
                 f"Profiler contents only supports {__support_contents}, but gets {content}"

verl/utils/profiler/torch_profile.py

@@ -13,16 +13,21 @@
 # limitations under the License.
 
 import functools
+import logging
 import os
 import re
 from datetime import datetime, timezone
 from typing import Callable, Optional
 
 import torch
 
+from verl.plugin.platform import get_platform
+
 from .config import ProfilerConfig, TorchProfilerToolConfig
 from .profile import DistProfiler
 
+logger = logging.getLogger(__name__)
+
 
 def get_dist_topology() -> dict:
     """Best-effort snapshot of the current process's distributed topology.
@@ -125,7 +130,9 @@ def get_torch_profiler(
         contents: Selects the other ``torch.profiler.profile`` arguments -- ``cuda`` maps to
             ``activities``, ``shapes`` to ``record_shapes``, ``memory`` to ``profile_memory`` and
             ``stack`` to ``with_stack``. CPU activity is always on, since verl's per-stage
-            ``record_function`` markers are CPU-side events.
+            ``record_function`` markers are CPU-side events. A plugin platform may also accept
+            its own content name here (see ``PlatformBase.torch_profiler_content_name``), mapped
+            to that platform's own ``torch.profiler.ProfilerActivity``.
         save_path: Directory to write chrome traces to.
         role: Optional logical scope name (e.g. ``train`` for a worker's whole-step window, or
             a stage name in discrete mode), embedded in the filename.
@@ -196,8 +203,27 @@ def _trace_handler(prof):
     # record_function, and those markers -- like operator names -- are CPU-side events, so a
     # device-only trace would be bare kernels that cannot be attributed to any stage.
     activities = [torch.profiler.ProfilerActivity.CPU]
-    if not contents or "cuda" in contents:
+    plugin_activity = get_platform().torch_profiler_activity()
+    plugin_content = get_platform().torch_profiler_content_name()
+    has_plugin_device = plugin_activity is not None and plugin_content is not None
+    plugin_requested = has_plugin_device and (not contents or plugin_content in contents)
+    # torch.profiler records a single device type and checks CUDA before every other device,
+    # so CUDA must stay out of `activities` whenever the plugin device was actually requested.
+    if not plugin_requested and (not contents or "cuda" in contents):
         activities.append(torch.profiler.ProfilerActivity.CUDA)
+    if plugin_requested:
+        activities.append(plugin_activity)
+        if "cuda" in contents:
+            # Only one device activity is ever recorded, and there is no active CUDA
+            # platform here (CUDA never implements these hooks) -- "cuda" is a dead
+            # keyword in this contents list, not a second activity.
+            logger.warning(
+                "profiler contents=%s requested both 'cuda' and '%s': only '%s' "
+                "(this platform's activity) is recorded, 'cuda' is ignored.",
+                sorted(contents),
+                plugin_content,
+                plugin_content,
+            )
 
     profile_kwargs = dict(
         activities=activities,