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 6cf4ece46..a557400b9 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 @@ -63,7 +63,7 @@ class ViTNpuGraphRunner(ViTCudaGraphRunner): def _create_graph( self, - graph_key: int, + graph_key: Hashable, ): graph = torch_npu.npu.NPUGraph() @@ -132,9 +132,16 @@ class ViTNpuGraphRunner(ViTCudaGraphRunner): cu_seqlens: torch.Tensor, rotary_pos_emb_cos: Optional[torch.Tensor] = None, rotary_pos_emb_sin: Optional[torch.Tensor] = None, - ) -> int: + attention_layout_key: Optional[Hashable] = None, + ) -> Hashable: vit = self.vit - graph_key = self._get_graph_key(x_3d) + graph_key = self._get_graph_key( + x_3d, + cu_seqlens, + None, + attention_layout_key, + ) + seq_len = x_3d.shape[0] if graph_key in self.block_graphs: return graph_key @@ -155,7 +162,7 @@ class ViTNpuGraphRunner(ViTCudaGraphRunner): self.block_output[graph_key] = x_3d self.block_input[graph_key] = x_3d self.block_ws[graph_key] = torch.empty( - graph_key, + seq_len, num_heads, attn_head_dim, device=self.device, @@ -178,7 +185,7 @@ class ViTNpuGraphRunner(ViTCudaGraphRunner): def replay( self, - graph_key: int, + graph_key: Hashable, x_3d: torch.Tensor, rotary_pos_emb_cos: Optional[torch.Tensor] = None, rotary_pos_emb_sin: Optional[torch.Tensor] = None, @@ -210,16 +217,23 @@ class ViTNpuGraphRunner(ViTCudaGraphRunner): rotary_pos_emb_cos: Optional[torch.Tensor] = None, rotary_pos_emb_sin: Optional[torch.Tensor] = None, output_indices: Optional[torch.Tensor] = None, + attention_layout_key: Optional[Hashable] = None, ) -> torch.Tensor: # x: [seq_len, hidden] -> [S, B=1, H] x_3d = x.unsqueeze(1) - graph_key = self._get_graph_key(x_3d) + graph_key = self._get_graph_key( + x_3d, + cu_seqlens, + None, + attention_layout_key, + ) if graph_key not in self.block_graphs: self.create_graph( x_3d=x_3d, cu_seqlens=cu_seqlens, rotary_pos_emb_cos=rotary_pos_emb_cos, rotary_pos_emb_sin=rotary_pos_emb_sin, + attention_layout_key=attention_layout_key, ) return self.replay( diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 5d2faac6f..f282eedbe 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -1029,6 +1029,7 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin): rotary_pos_emb_sin, ) = self._prepare_graph_inputs(x, grid_thw) + attention_layout_key = (tuple(cu_seqlens.tolist()), None) cu_seqlens = cu_seqlens.to("cpu") return self.graph_runners.run( x=x, @@ -1036,6 +1037,7 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin): rotary_pos_emb_sin=rotary_pos_emb_sin, cu_seqlens=cu_seqlens, output_indices=None, + attention_layout_key=attention_layout_key, ) def forward_with_cuda_graph( diff --git a/test/registered/unit/multimodal/test_vit_npu_graph_runner.py b/test/registered/unit/multimodal/test_vit_npu_graph_runner.py new file mode 100644 index 000000000..ebe2be344 --- /dev/null +++ b/test/registered/unit/multimodal/test_vit_npu_graph_runner.py @@ -0,0 +1,101 @@ +import importlib +import sys +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class _Block: + attn = SimpleNamespace( + qkv_backend_name="ascend_attn", + num_attention_heads_per_partition=2, + head_size=4, + ) + + def forward(self, x, output_ws=None): + return x + + +class _FakeGraph: + def replay(self): + pass + + +def _load_npu_graph_runner(): + # Load shared modules before stubbing torch_npu so platform detection stays on CPU. + importlib.import_module("sglang.srt.multimodal.vit_cuda_graph_runner") + torch_npu = SimpleNamespace() + with patch.dict( + sys.modules, + { + "torch_npu": torch_npu, + }, + ): + module = importlib.import_module( + "sglang.srt.hardware_backend.npu.graph_runner.vit_npu_graph_runner" + ) + return module.ViTNpuGraphRunner + + +def test_npu_vit_graph_keys_include_attention_boundaries(): + runner_cls = _load_npu_graph_runner() + vit = SimpleNamespace( + blocks=[_Block()], + merger=lambda x: x, + device=torch.device("cpu"), + dtype=torch.float32, + deepstack_visual_indexes=[], + deepstack_merger_list=None, + ) + + with patch( + "torch.get_device_module", + return_value=SimpleNamespace(graph_pool_handle=lambda: object()), + ): + runner = runner_cls(vit) + + runner_cls._graph_memory_pool = None + runner._create_graph = lambda graph_key: runner.block_graphs.__setitem__( + graph_key, _FakeGraph() + ) + + x = torch.zeros(8, 8) + rotary = torch.zeros(8, 4) + first_layout = torch.tensor([0, 4, 8], dtype=torch.int32) + second_layout = torch.tensor([0, 2, 8], dtype=torch.int32) + + with patch( + "sglang.srt.hardware_backend.npu.graph_runner." + "vit_npu_graph_runner.set_graph_pool_id" + ): + runner.run( + x, + first_layout, + rotary_pos_emb_cos=rotary, + rotary_pos_emb_sin=rotary, + ) + runner.run( + x, + second_layout, + rotary_pos_emb_cos=rotary, + rotary_pos_emb_sin=rotary, + ) + + assert len(runner.block_graphs) == 2 + assert {key[1][0] for key in runner.block_graphs} == { + (0, 4, 8), + (0, 2, 8), + } + assert all( + workspace.shape[0] == x.shape[0] for workspace in runner.block_ws.values() + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))