diff --git a/tests/data/test_valid_seqlens_devices.py b/tests/data/test_valid_seqlens_devices.py new file mode 100644 index 0000000000..0f2719496c --- /dev/null +++ b/tests/data/test_valid_seqlens_devices.py @@ -0,0 +1,73 @@ +# Copyright 2026 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. + +"""Tail-segment extraction must agree on CPU and Ascend NPU.""" + +import json +import os +from pathlib import Path + +import pytest +import torch + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.utils.device import IS_NPU_AVAILABLE +from veomni.utils.seqlen_pos_transform_utils import valid_seqlens_from_cu_seqlens + + +_CASES = [ + ([0, 4, 7, 8, 9, 10], None, [4, 3]), + ([0, 4, 5, 6, 7, 8], 3, [4, 1]), + ([0, 4, 5], 0, [4, 1]), + ([0, 1, 3], None, [1, 2]), + ([0, 1, 2], None, []), + ([0], None, []), +] + + +def _npu_test_eligible(): + """Honor the CI lease; on NPU CI leave device availability to the test assertion.""" + capabilities_file = os.environ.get("VEOMNI_CI_CAPABILITIES") + if capabilities_file: + return json.loads(Path(capabilities_file).read_text())["backend"] == "npu" + return IS_NPU_AVAILABLE and torch.npu.is_available() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) +@pytest.mark.parametrize(("backend", "expected"), [("cpu", False), ("mlu", False), ("npu", True)]) +def test_npu_eligibility_respects_ci_backend(monkeypatch, tmp_path, backend, expected): + capabilities = tmp_path / "capabilities.json" + capabilities.write_text(json.dumps({"backend": backend, "devices": []})) + monkeypatch.setenv("VEOMNI_CI_CAPABILITIES", str(capabilities)) + assert _npu_test_eligible() is expected + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) +@pytest.mark.parametrize("cu,tail,expected", _CASES) +def test_valid_seqlens_cpu(cu, tail, expected): + value = torch.tensor(cu, dtype=torch.int32) + result = valid_seqlens_from_cu_seqlens(value, tail_padding_length=tail) + assert result.tolist() == expected + assert result.dtype == torch.int32 + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="accelerator")) +@pytest.mark.skipif(not _npu_test_eligible(), reason="requires an Ascend NPU") +@pytest.mark.parametrize("cu,tail,expected", _CASES) +def test_valid_seqlens_ascend(cu, tail, expected): + assert IS_NPU_AVAILABLE and torch.npu.is_available(), "Ascend CI must provide an available NPU" + value = torch.tensor(cu, dtype=torch.int32, device="npu") + result = valid_seqlens_from_cu_seqlens(value, tail_padding_length=tail) + assert result.cpu().tolist() == expected + assert result.dtype == torch.int32 diff --git a/veomni/utils/seqlen_pos_transform_utils.py b/veomni/utils/seqlen_pos_transform_utils.py index 2d6adb58ae..7032528767 100644 --- a/veomni/utils/seqlen_pos_transform_utils.py +++ b/veomni/utils/seqlen_pos_transform_utils.py @@ -157,5 +157,5 @@ def valid_seqlens_from_cu_seqlens(cu_seqlens: torch.Tensor, tail_padding_length: valid_end = torch.clamp(cu_seqlens[-1:] - tail_padding_length, min=0) return diff[cu_seqlens[1:] <= valid_end] - pad = int((torch.flip(diff == 1, (0,)).cumprod(0)).sum().item()) + pad = int((torch.flip(diff == 1, (0,)).int().cumprod(0)).sum().item()) return diff[:-pad] if pad else diff