[hardware, perf, training_utils] feat: add a profiler_markers() platform hook and share the attention padding helpers

pyproject.toml

@@ -67,6 +67,7 @@ verl-core = [
     # Decoding needs FFmpeg's shared libs (`ffmpeg`, docker/Dockerfile.uv.cu130).
     "torchcodec",
     "cachetools",
+    "einops",
     "nvtx",
     "pytest",
     "pytest-asyncio",

tests/models/test_transformer.py

@@ -23,13 +23,8 @@
     Qwen2Config,
 )
 
+from verl.utils.attention_utils import index_first_axis, pad_input, rearrange, unpad_input
 from verl.utils.device import get_device_name
-
-if get_device_name() == "cuda":
-    from flash_attn.bert_padding import index_first_axis, pad_input, rearrange, unpad_input
-elif get_device_name() == "npu":
-    from verl.utils.attention_utils import index_first_axis, pad_input, rearrange, unpad_input
-
 from verl.utils.model import compute_position_id_with_mask, create_random_mask
 from verl.utils.torch_functional import log_probs_from_logits_all_rmpad, masked_mean
 

tests/models/test_transformers_ulysses.py

@@ -25,6 +25,7 @@
 
 from verl.models.transformers.monkey_patch import apply_monkey_patch
 from verl.protocol import DataProto
+from verl.utils.attention_utils import index_first_axis, rearrange, unpad_input
 from verl.utils.device import get_device_name, get_torch_device
 from verl.utils.distributed import initialize_global_process_group
 from verl.utils.model import compute_position_id_with_mask, create_random_mask
@@ -36,11 +37,6 @@
     ulysses_pad_and_slice_inputs,
 )
 
-if get_device_name() == "cuda":
-    from flash_attn.bert_padding import index_first_axis, rearrange, unpad_input
-elif get_device_name() == "npu":
-    from verl.utils.attention_utils import index_first_axis, rearrange, unpad_input
-
 # TODO(sgm): add more models for test
 # we only need one scale for each model
 

tests/plugin/test_platform_abstraction.py

@@ -131,7 +131,7 @@ def test_case_insensitive(self):
 
     def test_empty_triggers_auto_detection(self):
         with mock.patch.dict(os.environ, {"VERL_PLATFORM": ""}):
-            assert _detect_platform_name() in ("nvidia", "huawei", "amd")
+            assert _detect_platform_name() in ("nvidia", "huawei", "amd", "intel")
 
 
 class TestPlatformCreation:

tests/special_distributed/test_fsdp_ckpt.py

@@ -22,17 +22,14 @@
 from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
 from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2Config
 
+from verl.utils.attention_utils import unpad_input
 from verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager
 from verl.utils.device import get_device_name, get_torch_device
 from verl.utils.distributed import initialize_global_process_group
 from verl.utils.fsdp_utils import MixedPrecisionPolicy, apply_fsdp2
 
 
 def create_random_input_ids(batch_size, seq_len, vocab_size):
-    if get_device_name() == "cuda":
-        from flash_attn.bert_padding import unpad_input
-    elif get_device_name() == "npu":
-        from verl.utils.attention_utils import unpad_input
     from verl.utils.model import compute_position_id_with_mask, create_random_mask
 
     input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=get_device_name())

tests/utils/test_activation_offload.py

@@ -25,16 +25,13 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2Config
 
 from verl.utils.activation_offload import enable_activation_offloading
+from verl.utils.attention_utils import unpad_input
 from verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager
 from verl.utils.device import get_device_name, get_nccl_backend, get_torch_device
 from verl.utils.fsdp_utils import MixedPrecisionPolicy, apply_fsdp2, get_fsdp_wrap_policy
 
 
 def create_random_input_ids(batch_size, seq_len, vocab_size):
-    if get_device_name() == "cuda":
-        from flash_attn.bert_padding import unpad_input
-    elif get_device_name() == "npu":
-        from verl.utils.attention_utils import unpad_input
     from verl.utils.model import compute_position_id_with_mask, create_random_mask
 
     input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=get_device_name())

tests/utils/test_attention_utils.py

@@ -0,0 +1,115 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+"""Unit tests for the backend dispatch in verl.utils.attention_utils."""
+
+import sys
+import types
+from unittest import mock
+
+import pytest
+import torch
+
+from verl.utils import attention_padding_utils, attention_utils
+
+
+@pytest.fixture
+def reset_resolved_functions():
+    """The resolved functions live in module globals; clear them around each test."""
+    attention_utils._index_first_axis = None
+    attention_utils._pad_input = None
+    attention_utils._rearrange = None
+    attention_utils._unpad_input = None
+    yield
+    attention_utils._index_first_axis = None
+    attention_utils._pad_input = None
+    attention_utils._rearrange = None
+    attention_utils._unpad_input = None
+
+
+def _fake_flash_attn():
+    """A stand-in for the flash_attn package, which has no wheel on non-CUDA devices."""
+    pkg = types.ModuleType("flash_attn")
+    bert_padding = types.ModuleType("flash_attn.bert_padding")
+    for name in ("index_first_axis", "pad_input", "rearrange", "unpad_input"):
+        setattr(bert_padding, name, mock.Mock(name=name, return_value=f"flash_{name}"))
+    pkg.bert_padding = bert_padding
+    return {"flash_attn": pkg, "flash_attn.bert_padding": bert_padding}
+
+
+def test_prefers_flash_attn_when_installed(reset_resolved_functions):
+    """With flash-attn importable, all four entry points dispatch to it."""
+    with (
+        mock.patch.dict(sys.modules, _fake_flash_attn()),
+        mock.patch("verl.utils.device.is_torch_npu_available", return_value=False),
+    ):
+        assert attention_utils.index_first_axis() == "flash_index_first_axis"
+        assert attention_utils.pad_input() == "flash_pad_input"
+        assert attention_utils.rearrange() == "flash_rearrange"
+        assert attention_utils.unpad_input() == "flash_unpad_input"
+
+
+def test_falls_back_to_padding_utils_without_flash_attn(reset_resolved_functions):
+    """Without flash-attn, dispatch goes to the pure-torch port rather than a local duplicate."""
+    # A None entry in sys.modules makes `import flash_attn...` raise ImportError.
+    with (
+        mock.patch.dict(sys.modules, {"flash_attn": None, "flash_attn.bert_padding": None}),
+        mock.patch("verl.utils.device.is_torch_npu_available", return_value=False),
+    ):
+        index_first_axis, pad_input, rearrange, unpad_input = attention_utils._get_attention_functions()
+
+    assert index_first_axis is attention_padding_utils.index_first_axis
+    assert pad_input is attention_padding_utils.pad_input
+    assert rearrange is attention_padding_utils.rearrange
+    assert unpad_input is attention_padding_utils.unpad_input
+
+
+def test_npu_uses_padding_utils_over_flash_attn(reset_resolved_functions):
+    """NPU keeps its existing behaviour: the port wins even if flash-attn is importable."""
+    with (
+        mock.patch.dict(sys.modules, _fake_flash_attn()),
+        mock.patch("verl.utils.device.is_torch_npu_available", return_value=True),
+    ):
+        index_first_axis, *_ = attention_utils._get_attention_functions()
+
+    assert index_first_axis is attention_padding_utils.index_first_axis
+
+
+def test_padding_utils_unpad_pad_round_trip():
+    """pad_input(unpad_input(x)) must restore the padded tensor and report correct metadata."""
+    hidden_states = torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4)
+    attention_mask = torch.tensor([[1, 1, 0], [1, 0, 0]])
+
+    unpadded, indices, cu_seqlens, max_seqlen, seqused = attention_padding_utils.unpad_input(
+        hidden_states, attention_mask
+    )
+
+    assert unpadded.shape == (3, 4)
+    assert max_seqlen == 2
+    torch.testing.assert_close(cu_seqlens, torch.tensor([0, 2, 3], dtype=torch.int32))
+    torch.testing.assert_close(seqused, torch.tensor([2, 1], dtype=torch.int32))
+
+    repadded = attention_padding_utils.pad_input(unpadded, indices, batch=2, seqlen=3)
+    torch.testing.assert_close(repadded, hidden_states * attention_mask.unsqueeze(-1))
+
+
+def test_padding_utils_index_first_axis_is_differentiable():
+    """index_first_axis is an autograd.Function; gradients must flow back to the unselected rows."""
+    tensor = torch.randn(4, 3, requires_grad=True)
+    indices = torch.tensor([0, 2])
+
+    attention_padding_utils.index_first_axis(tensor, indices).sum().backward()
+
+    expected = torch.zeros(4, 3)
+    expected[indices] = 1.0
+    torch.testing.assert_close(tensor.grad, expected)

tests/utils/test_profiler_markers.py

@@ -0,0 +1,129 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+"""Unit tests for lazy profiler-marker resolution (verl.utils.profiler._resolve_markers)."""
+
+from unittest import mock
+
+import pytest
+
+import verl.plugin.platform.platform_manager as pm
+import verl.utils.profiler as profiler_pkg
+from verl.plugin.platform import set_platform
+from verl.utils.profiler import mark_annotate, mark_end_range, mark_start_range, marked_timer
+
+
+@pytest.fixture
+def reset_profiler_and_platform():
+    def _clear():
+        pm._current_platform = None
+        profiler_pkg._mark_start_range = None
+        profiler_pkg._mark_end_range = None
+        profiler_pkg._mark_annotate = None
+        profiler_pkg._marked_timer = None
+
+    _clear()
+    yield
+    _clear()
+
+
+def _platform_with_markers():
+    """A platform supplying all four markers, plus the mocks it returns."""
+    markers = (
+        mock.Mock(return_value="range-id"),
+        mock.Mock(),
+        mock.Mock(side_effect=lambda *a, **k: lambda f: f),
+        mock.Mock(),
+    )
+    platform = mock.Mock()
+    platform.profiler_markers.return_value = markers
+    return platform, markers
+
+
+def test_platform_markers_win_over_an_importable_nvtx(reset_profiler_and_platform):
+    """nvtx ships in the verl-core extra, so is_nvtx_available() can be True on a device that
+    cannot use it. A platform that supplies its own markers must not be shadowed by that."""
+    platform, (plugin_start, plugin_end, plugin_annotate, plugin_timer) = _platform_with_markers()
+    set_platform(platform)
+
+    with mock.patch("verl.utils.profiler.is_nvtx_available", return_value=True):
+        # Not resolved until first use.
+        assert profiler_pkg._mark_start_range is None
+
+        result = mark_start_range("hello")
+        mark_end_range("range-id")
+        mark_annotate()(lambda: None)
+        marked_timer()
+
+    assert result == "range-id"
+    plugin_start.assert_called_once_with("hello")
+    plugin_end.assert_called_once_with("range-id")
+    plugin_annotate.assert_called_once()
+    plugin_timer.assert_called_once()
+    # Resolved once, then cached for subsequent calls.
+    assert profiler_pkg._mark_start_range is plugin_start
+    platform.profiler_markers.assert_called_once()
+
+
+def test_nvtx_still_used_when_platform_supplies_no_markers(reset_profiler_and_platform):
+    """CUDA/NPU behaviour is unchanged: PlatformBase.profiler_markers() returns None there."""
+    pytest.importorskip("nvtx", reason="nvtx_profile imports the nvtx package at module level")
+
+    platform = mock.Mock()
+    platform.profiler_markers.return_value = None
+    set_platform(platform)
+
+    with mock.patch("verl.utils.profiler.is_nvtx_available", return_value=True):
+        profiler_pkg._resolve_markers()
+
+    from verl.utils.profiler import nvtx_profile
+
+    assert profiler_pkg._mark_start_range is nvtx_profile.mark_start_range
+    assert profiler_pkg._marked_timer is nvtx_profile.marked_timer
+
+
+def test_npu_falls_through_to_mstx(reset_profiler_and_platform):
+    """NPU without nvtx keeps using mstx markers."""
+    pytest.importorskip("torch_npu", reason="mstx_profile imports torch_npu at module level")
+
+    platform = mock.Mock()
+    platform.profiler_markers.return_value = None
+    set_platform(platform)
+
+    with (
+        mock.patch("verl.utils.profiler.is_nvtx_available", return_value=False),
+        mock.patch("verl.utils.profiler.is_npu_available", True),
+    ):
+        profiler_pkg._resolve_markers()
+
+    from verl.utils.profiler import mstx_profile
+
+    assert profiler_pkg._mark_start_range is mstx_profile.mark_start_range
+
+
+def test_generic_fallback_when_nothing_applies(reset_profiler_and_platform):
+    """No platform markers, no nvtx, no npu -> the pure-Python no-op markers."""
+    platform = mock.Mock()
+    platform.profiler_markers.return_value = None
+    set_platform(platform)
+
+    with (
+        mock.patch("verl.utils.profiler.is_nvtx_available", return_value=False),
+        mock.patch("verl.utils.profiler.is_npu_available", False),
+    ):
+        profiler_pkg._resolve_markers()
+
+    from verl.utils.profiler import performance, profile
+
+    assert profiler_pkg._mark_start_range is profile.mark_start_range
+    assert profiler_pkg._marked_timer is performance.marked_timer