[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>

python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py

@@ -27,7 +27,6 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
 )
 from sglang.srt.layers.attention.vision import VisionAttention
 from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
-from sglang.srt.runtime_context import get_mm
 
 
 class ViTNpuGraphRunner(ViTCudaGraphRunner):
@@ -70,17 +69,19 @@ class ViTNpuGraphRunner(ViTCudaGraphRunner):
         graph = torch_npu.npu.NPUGraph()
         vit = self.vit
 
-        override_backend = get_mm().mm_attention_backend
+        backend = self._attn_backend
         with torch_npu.npu.graph(graph, pool=ViTNpuGraphRunner._graph_memory_pool):
             y = None
             deepstack_outs: List[torch.Tensor] = []
             deepstack_capture_idx = 0
 
             for layer_num, blk in enumerate(vit.blocks):
-                if override_backend == "ascend_attn":
+                if backend == "ascend_attn":
                     cu_seq_lens = self.cu_seq_lens[graph_key]
                 else:
-                    raise RuntimeError("Not supported ViT attention backend")
+                    raise RuntimeError(
+                        f"ViT NPU graph does not support attention backend: {backend}"
+                    )
 
                 if layer_num == 0:
                     y = blk(

python/sglang/srt/layers/attention/vision.py

@@ -784,6 +784,9 @@ class VisionAscendAttention(nn.Module):
         if not _is_npu:
             raise Exception("VisionAscendAttention is only available for ascend npu")
         super().__init__()
+        # Ascend fused attention does not support SGLang's additive masks, so
+        # masked inputs must stay on SDPA.
+        self.sdpa_fallback = VisionSdpaAttention(**kwargs)
 
     def forward(
         self,
@@ -795,6 +798,7 @@ class VisionAscendAttention(nn.Module):
         seq_len: int,
         softmax_scale: Optional[float] = None,
         forward_metadata: Optional[VisionAttentionMetadata] = None,
+        attention_mask: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
         r"""
@@ -803,6 +807,19 @@ class VisionAscendAttention(nn.Module):
         Returns:
              [b * s, h, head_size]
         """
+        if attention_mask is not None:
+            return self.sdpa_fallback(
+                q=q,
+                k=k,
+                v=v,
+                cu_seqlens=cu_seqlens,
+                bsz=bsz,
+                seq_len=seq_len,
+                attention_mask=attention_mask,
+                forward_metadata=forward_metadata,
+                **kwargs,
+            )
+
         if forward_metadata is not None:
             # TND fused attention expects cumulative seqlens (cu_seqlens[1:]),
             # not per-sequence lengths in forward_metadata.seq_lens.
@@ -1050,6 +1067,8 @@ class VisionAttention(nn.Module):
             print_info_once(f"Multimodal attention backend not set. Use {qkv_backend}.")
         print_info_once(f"Using {qkv_backend} as multimodal attention backend.")
 
+        self.qkv_backend_name: str = qkv_backend
+
         self.customized_position_embedding_applier = (
             customized_position_embedding_applier
         )
@@ -1151,7 +1170,8 @@ class VisionAttention(nn.Module):
         - CUDA (Hopper SM90): "fa3"
         - CUDA (Blackwell SM100): "fa4"
         - CUDA (other): "triton_attn"
-        - Non-CUDA: "sdpa"
+        - Ascend NPU: "ascend_attn"
+        - Other platforms: device-specific optimized backend or "sdpa"
         """
         override_backend = get_mm().mm_attention_backend
         if override_backend is not None:
@@ -1166,6 +1186,8 @@ class VisionAttention(nn.Module):
                 backend = "fa4"
             else:
                 backend = "triton_attn"
+        elif _is_npu:
+            backend = "ascend_attn"
         elif _is_musa:
             if get_device_capability() >= (3, 1):
                 backend = "fa3"

python/sglang/srt/multimodal/internvl_vit_cuda_graph_runner.py

@@ -22,7 +22,6 @@ import torch
 import torch.nn as nn
 
 from sglang.srt.layers.attention.vision import VisionAttention
-from sglang.srt.runtime_context import get_mm
 
 
 class InternViTCudaGraphRunner:
@@ -51,6 +50,7 @@ class InternViTCudaGraphRunner:
         first_layer = encoder.layers[0]
         # InternAttention wraps VisionAttention as first_layer.attn.attn
         self._attn: VisionAttention = first_layer.attn.attn  # type: ignore
+        self._attn_backend: str | None = getattr(self._attn, "qkv_backend_name", None)
 
     @property
     def device(self) -> torch.device:
@@ -95,17 +95,19 @@ class InternViTCudaGraphRunner:
 
     def _warmup_once(self, key: Hashable) -> None:
         """Run a tiny eager warmup on the preallocated buffers to trigger lazy init."""
-        override_backend = get_mm().mm_attention_backend
+        backend = self._attn_backend
         cu = self.cu[key]
         cu_kk = self.cu_kk[key]
         max_len = int(cu_kk.max().item()) if cu_kk.numel() else 0
 
-        if override_backend == "triton_attn":
+        if backend == "triton_attn":
             cu_ws = [cu, cu_kk, max_len]
-        elif override_backend == "fa3":
+        elif backend == "fa3":
             cu_ws = [cu, max_len]
         else:
-            raise RuntimeError("Not supported ViT attention backend for InternVL CG")
+            raise RuntimeError(
+                f"InternVL ViT CUDA graph does not support attention backend: {backend}"
+            )
 
         x = self.inp[key]
         y = x
@@ -115,18 +117,20 @@ class InternViTCudaGraphRunner:
 
     def _capture_graph(self, key: Hashable) -> None:
         g = torch.cuda.CUDAGraph()
-        override_backend = get_mm().mm_attention_backend
+        backend = self._attn_backend
 
         cu = self.cu[key]
         cu_kk = self.cu_kk[key]
         max_len = int(cu_kk.max().item()) if cu_kk.numel() else 0
 
-        if override_backend == "triton_attn":
+        if backend == "triton_attn":
             cu_ws = [cu, cu_kk, max_len]
-        elif override_backend == "fa3":
+        elif backend == "fa3":
             cu_ws = [cu, max_len]
         else:
-            raise RuntimeError("Not supported ViT attention backend for InternVL CG")
+            raise RuntimeError(
+                f"InternVL ViT CUDA graph does not support attention backend: {backend}"
+            )
 
         torch.cuda.synchronize()
 

python/sglang/srt/multimodal/vit_cuda_graph_runner.py

@@ -25,7 +25,6 @@ import torch.nn as nn
 
 from sglang.srt.distributed.parallel_state import get_tp_group
 from sglang.srt.layers.attention.vision import VisionAttention
-from sglang.srt.runtime_context import get_mm
 
 
 class ViTCudaGraphRunner:
@@ -80,7 +79,9 @@ class ViTCudaGraphRunner:
         )
 
         self._attn: Optional[VisionAttention] = getattr(first_blk, "attn", None)
-        self._attn_backend = getattr(self._attn, "qkv_backend", None)
+        self._attn_backend: Optional[str] = getattr(
+            self._attn, "qkv_backend_name", None
+        )
 
     @property
     def device(self) -> torch.device:
@@ -151,13 +152,13 @@ class ViTCudaGraphRunner:
         cu_full_kk = self.cu_full_len_kk[graph_key]
         max_full_len = int(cu_full_kk.max().item())
 
-        override_backend = get_mm().mm_attention_backend
+        backend = self._attn_backend
 
         if self._fullatt_block_indexes and 0 not in vit.fullatt_block_indexes:
             warmup_cu_ws = [cu_window, cu_window_kk, max_window_len]
         else:
             warmup_cu_ws = [cu_full, cu_full_kk, max_full_len]
-        if override_backend == "fa3":
+        if backend == "fa3":
             warmup_cu_ws = [warmup_cu_ws[0], warmup_cu_ws[2]]
 
         warmup_kwargs = dict(
@@ -192,12 +193,14 @@ class ViTCudaGraphRunner:
                     cu_seqlens_kk_now = cu_full_kk
                     max_len = max_full_len
 
-                if override_backend == "triton_attn":
+                if backend == "triton_attn":
                     cu_seq_len_ws = [cu_seqlens_now, cu_seqlens_kk_now, max_len]
-                elif override_backend == "fa3":
+                elif backend == "fa3":
                     cu_seq_len_ws = [cu_seqlens_now, max_len]
                 else:
-                    raise RuntimeError("Not supported ViT attention backend")
+                    raise RuntimeError(
+                        f"ViT CUDA graph does not support attention backend: {backend}"
+                    )
 
                 if position_embeddings is not None:
                     if layer_num == 0:

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__]))