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