[NPU] Enable automatic ascend_attn selection for vision attention and graph runners (#31948)

Co-authored-by: litao.dream <litao.dream@bytedance.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>

test/registered/unit/layers/attention/test_vision_backend_selection.py

@@ -0,0 +1,151 @@
+import sys
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+from sglang.srt.layers.attention import vision
+from sglang.test.ci.ci_register import register_cpu_ci, register_npu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+register_npu_ci(est_time=2, suite="stage-b-test-1-npu-a2")
+
+
+@pytest.fixture
+def npu_platform(monkeypatch):
+    monkeypatch.setattr(vision, "is_cuda", lambda: False)
+    monkeypatch.setattr(vision, "_is_npu", True)
+    monkeypatch.setattr(vision, "_is_musa", False)
+    monkeypatch.setattr(vision, "_is_hip", False)
+    monkeypatch.setattr(vision, "_is_cpu", False)
+    monkeypatch.setattr(vision, "_is_xpu", False)
+
+
+@pytest.mark.parametrize(
+    ("server_backend", "passed_backend", "expected"),
+    [
+        (None, None, "ascend_attn"),
+        (None, "sdpa", "sdpa"),
+        ("sdpa", None, "sdpa"),
+        ("sdpa", "ascend_attn", "sdpa"),
+    ],
+)
+def test_npu_backend_selection_priority(
+    monkeypatch,
+    npu_platform,
+    server_backend,
+    passed_backend,
+    expected,
+):
+    monkeypatch.setattr(
+        vision,
+        "get_mm",
+        lambda: SimpleNamespace(mm_attention_backend=server_backend),
+    )
+
+    backend = vision.VisionAttention._determine_attention_backend(None, passed_backend)
+
+    assert backend == expected
+
+
+@pytest.mark.parametrize("mask_kind", ["causal", "padding"])
+def test_ascend_attention_masked_inputs_fall_back_to_sdpa(
+    monkeypatch,
+    npu_platform,
+    mask_kind,
+):
+    torch.manual_seed(0)
+    bsz, seq_len, num_heads, head_dim = 2, 4, 2, 8
+    softmax_scale = 0.37
+    q, k, v = [torch.randn(bsz * seq_len, num_heads, head_dim) for _ in range(3)]
+    mask = torch.zeros(bsz, 1, seq_len, seq_len)
+    if mask_kind == "causal":
+        masked_positions = torch.ones(seq_len, seq_len, dtype=torch.bool).triu(1)
+        mask.masked_fill_(masked_positions, torch.finfo(mask.dtype).min)
+    else:
+        mask[:, :, :, -1] = torch.finfo(mask.dtype).min
+
+    fused_attention = Mock(
+        side_effect=AssertionError("masked inputs must not use Ascend fused attention")
+    )
+    monkeypatch.setattr(
+        vision,
+        "torch_npu",
+        SimpleNamespace(npu_fused_infer_attention_score=fused_attention),
+        raising=False,
+    )
+    backend = vision.VisionAscendAttention(
+        head_dim=head_dim,
+        num_heads=num_heads,
+        num_kv_heads=num_heads,
+        softmax_scale=softmax_scale,
+    )
+
+    output = backend(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens=torch.arange(0, (bsz + 1) * seq_len, seq_len),
+        bsz=bsz,
+        seq_len=seq_len,
+        attention_mask=mask,
+    )
+    q_ref, k_ref, v_ref = [
+        rearrange(x, "(b s) h d -> b h s d", b=bsz) for x in (q, k, v)
+    ]
+    expected = F.scaled_dot_product_attention(
+        q_ref,
+        k_ref,
+        v_ref,
+        attn_mask=mask,
+        scale=softmax_scale,
+    )
+    expected = rearrange(expected, "b h s d -> (b s) h d")
+
+    torch.testing.assert_close(output, expected)
+    fused_attention.assert_not_called()
+
+
+def test_ascend_attention_unmasked_inputs_keep_fused_path(
+    monkeypatch,
+    npu_platform,
+):
+    bsz, seq_len, num_heads, head_dim = 1, 2, 2, 8
+    q, k, v = [torch.randn(bsz * seq_len, num_heads, head_dim) for _ in range(3)]
+    expected = torch.randn_like(q)
+    fused_attention = Mock(return_value=(expected, None))
+    monkeypatch.setattr(
+        vision,
+        "torch_npu",
+        SimpleNamespace(npu_fused_infer_attention_score=fused_attention),
+        raising=False,
+    )
+    backend = vision.VisionAscendAttention(
+        head_dim=head_dim,
+        num_heads=num_heads,
+        num_kv_heads=num_heads,
+    )
+    sdpa_forward = Mock(
+        side_effect=AssertionError("unmasked inputs must keep Ascend fused attention")
+    )
+    monkeypatch.setattr(backend.sdpa_fallback, "forward", sdpa_forward)
+
+    output = backend(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens=torch.tensor([0, seq_len], dtype=torch.int32),
+        bsz=bsz,
+        seq_len=seq_len,
+    )
+
+    torch.testing.assert_close(output, expected)
+    fused_attention.assert_called_once()
+    sdpa_forward.assert_not_called()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

test/registered/unit/multimodal/test_vit_cuda_graph_runner.py

@@ -1,3 +1,4 @@
+import sys
 from types import SimpleNamespace
 from unittest.mock import patch
 
@@ -7,6 +8,9 @@ from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=3, suite="base-a-test-cpu")
 
+from sglang.srt.multimodal.internvl_vit_cuda_graph_runner import (
+    InternViTCudaGraphRunner,
+)
 from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
 
 
@@ -55,5 +59,35 @@ def test_non_dp_vit_graph_capture_uses_tp_communication_capture():
     assert entered == [True]
 
 
+def test_vit_graph_runner_caches_resolved_backend_name():
+    class Block:
+        attn = SimpleNamespace(
+            qkv_backend_name="fa3",
+            qkv_backend=object(),
+        )
+
+        def forward(self, x, output_ws=None):
+            return x
+
+    vit = SimpleNamespace(blocks=[Block()])
+
+    runner = ViTCudaGraphRunner(vit)
+
+    assert runner._attn_backend == "fa3"
+
+
+def test_internvl_graph_runner_caches_resolved_backend_name():
+    attention = SimpleNamespace(
+        qkv_backend_name="triton_attn",
+        qkv_backend=object(),
+    )
+    layer = SimpleNamespace(attn=SimpleNamespace(attn=attention))
+    encoder = SimpleNamespace(layers=[layer])
+
+    runner = InternViTCudaGraphRunner(encoder)
+
+    assert runner._attn_backend == "triton_attn"
+
+
 if __name__ == "__main__":
-    raise SystemExit(pytest.main([__file__, "-v"]))
+    sys.exit(pytest.main([__file__]))