[NPU] Fix ViT graph key layout handling (#37758)

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(

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(

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