Add numerical tests for Ascend ViT varlen attention

tests/npu_ops/test_vit_attn.py

@@ -0,0 +1,107 @@
+"""Numerical coverage for the standalone Ascend ViT attention shim."""
+
+from __future__ import annotations
+
+import importlib.util
+import itertools
+import math
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+import pytest
+
+pytestmark = pytest.mark.ci_npu
+
+
+@pytest.fixture(scope="module")
+def npu_modules() -> tuple[ModuleType, ModuleType]:
+    torch = pytest.importorskip("torch")
+    pytest.importorskip("torch_npu")
+    if not torch.npu.is_available():
+        pytest.skip("requires an available Ascend NPU")
+
+    # Import the shim itself without installing npu_ops' process-wide import redirects.
+    path = Path(__file__).resolve().parents[2] / "deepinfer/npu_ops/vit_attn.py"
+    spec = importlib.util.spec_from_file_location("_deepinfer_vit_attn_test", path)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return torch, module
+
+
+@pytest.mark.parametrize("dtype_name", ["float16", "bfloat16"])
+@pytest.mark.parametrize("packed", [False, True], ids=["contiguous", "packed_qkv"])
+@pytest.mark.parametrize("lengths", [(1,), (17,), (3, 31, 65), (128, 127)])
+@pytest.mark.parametrize(("heads", "head_dim"), [(4, 64), (8, 128)])
+def test_varlen_attention_matches_reference(
+    npu_modules: tuple[ModuleType, ModuleType], dtype_name: str, packed: bool, lengths: tuple[int, ...], heads: int, head_dim: int
+) -> None:
+    torch, vit_attn = npu_modules
+    dtype = getattr(torch, dtype_name)
+    generator = torch.Generator().manual_seed(0)
+    data = (torch.randn(sum(lengths), 3, heads, head_dim, generator=generator) * 0.5).to(dtype)
+    cpu_q, cpu_k, cpu_v = data.unbind(1)
+    if packed:
+        # This is the layout produced by ViT's packed QKV projection and unbind.
+        q, k, v = data.to("npu:0").unbind(1)
+        assert q.stride(0) == 3 * heads * head_dim
+    else:
+        q, k, v = (tensor.contiguous().to("npu:0") for tensor in (cpu_q, cpu_k, cpu_v))
+
+    ends = list(itertools.accumulate(lengths))
+    output, lse = vit_attn._flash_attn_fwd(
+        q,
+        k,
+        v,
+        cu_seqlens_q=ends,
+        cu_seqlens_k=ends,
+        max_seqlen_q=max(lengths),
+        max_seqlen_k=max(lengths),
+        softmax_scale=None,
+        causal=False,
+        window_size_left=-1,
+        window_size_right=-1,
+        num_splits=1,
+        pack_gqa=None,
+        return_lse=False,
+    )
+
+    # Compute each segment independently in FP64 from the dtype-rounded inputs.
+    reference = []
+    start = 0
+    for end in ends:
+        q64, k64, v64 = (tensor[start:end].transpose(0, 1).double() for tensor in (cpu_q, cpu_k, cpu_v))
+        probabilities = ((q64 @ k64.mT) / math.sqrt(head_dim)).softmax(dim=-1)
+        reference.append((probabilities @ v64).transpose(0, 1))
+        start = end
+
+    assert lse is None
+    assert output.shape == q.shape
+    assert output.dtype == dtype
+    assert output.device == q.device
+    _assert_close(torch, output, torch.cat(reference), dtype_name)
+
+
+@pytest.mark.parametrize("dtype_name", ["float16", "bfloat16"])
+def test_varlen_attention_does_not_mix_segments(npu_modules: tuple[ModuleType, ModuleType], dtype_name: str) -> None:
+    torch, vit_attn = npu_modules
+    lengths = (3, 5, 1)
+    q = torch.zeros(sum(lengths), 4, 64, dtype=getattr(torch, dtype_name))
+    values = q.clone()
+    start = 0
+    for length, value in zip(lengths, (2, -3, 7), strict=True):
+        values[start : start + length].fill_(value)
+        start += length
+
+    # Uniform attention must preserve the distinct constant V of each segment.
+    output, lse = vit_attn._flash_attn_fwd(
+        q.to("npu:0"), q.to("npu:0"), values.to("npu:0"), cu_seqlens_q=list(itertools.accumulate(lengths))
+    )
+    assert lse is None
+    _assert_close(torch, output, values, dtype_name)
+
+
+def _assert_close(torch: ModuleType, actual: Any, expected: Any, dtype_name: str) -> None:
+    rtol, atol = (2e-2, 2e-3) if dtype_name == "bfloat16" else (3e-3, 3e-4)
+    torch.testing.assert_close(actual.float().cpu(), expected.float(), rtol=rtol, atol=atol)