diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py index 46dd3a101..6cf4ece46 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py +++ b/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( diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 25d065228..a6fcc0e37 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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" diff --git a/python/sglang/srt/multimodal/internvl_vit_cuda_graph_runner.py b/python/sglang/srt/multimodal/internvl_vit_cuda_graph_runner.py index 23ce493eb..1f60d8cf0 100644 --- a/python/sglang/srt/multimodal/internvl_vit_cuda_graph_runner.py +++ b/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() diff --git a/python/sglang/srt/multimodal/vit_cuda_graph_runner.py b/python/sglang/srt/multimodal/vit_cuda_graph_runner.py index 336a09b6a..fd6b5d246 100644 --- a/python/sglang/srt/multimodal/vit_cuda_graph_runner.py +++ b/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: diff --git a/test/registered/unit/layers/attention/test_vision_backend_selection.py b/test/registered/unit/layers/attention/test_vision_backend_selection.py new file mode 100644 index 000000000..50f241a55 --- /dev/null +++ b/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__])) diff --git a/test/registered/unit/multimodal/test_vit_cuda_graph_runner.py b/test/registered/unit/multimodal/test_vit_cuda_graph_runner.py index 7a7df286c..4d6c15a2f 100644 --- a/test/registered/unit/multimodal/test_vit_cuda_graph_runner.py +++ b/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__]))