[Feat] [AMD] vision: Support ViT CUDA graph with the aiter_attn backend

docs/docs/advanced_features/cuda_graph_for_multi_modal_encoder.mdx

@@ -43,12 +43,15 @@ Attention backend arguments are fixed inside the graph:
 
 TritonAttn expects [cu_seqlens, cu_seqlens_kk, max_len]
 FA3 expects [cu_seqlens, max_len]
+FA4 and AiterAttn (`aiter_attn`, ROCm) take the same [cu_seqlens, max_len] as FA3
 
 max_len is frozen as an int constant.
 cu_seqlens is cached into a dict during create_graph(), and its contents are not updated during subsequent replays.
 
 For the same graph_key = S, you not only require the input shape to match, but also require the segmentation pattern in cu_seqlens (and window seqlens) to be identical. Otherwise, attention will segment the sequence incorrectly.
 
+Because the key includes the segmentation pattern, every new image size or batch composition captures a new graph. At most `SGLANG_VIT_CUDA_GRAPH_CACHE_CAPACITY` graphs are captured (default 32); later layouts run the same blocks eagerly instead of growing the cache. Three image sizes with up to 8 images per ViT call need 18 graphs per rank, and each graph adds only its own input, output and attention-workspace buffers (about 7 KB per patch for a 1152-wide encoder); activations share one memory pool.
+
 ### Rotary buffer management
 The feature reallocates a larger sin_cos_ws when seq_len increases.
 The max_content_len is used to make sure the maximum size of the allocated rotary buffer.

docs/docs/references/environment_variables.mdx

@@ -2082,6 +2082,11 @@ MPS and does not read them.
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Capture the vision encoder (ViT) in a CUDA graph.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_VIT_CUDA_GRAPH_CACHE_CAPACITY</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Maximum number of ViT CUDA graphs (one per attention layout) to capture. Later layouts run the same code without a graph. Default 32 (three image sizes with up to 8 images per ViT call need 18 per rank).</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>32</td>
+    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_USE_CUDA_IPC_TRANSPORT</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use CUDA IPC transport for multimodal-item tensors.</td>

python/sglang/srt/environ.py

@@ -1599,6 +1599,10 @@ class Envs:
     SGLANG_MM_BUFFER_SIZE_MB = EnvInt(0)
     SGLANG_MM_PRECOMPUTE_HASH = EnvBool(False)
     SGLANG_VIT_ENABLE_CUDA_GRAPH = EnvBool(False)
+    # Max ViT CUDA graphs (one per attention layout) to capture; later layouts run
+    # without a graph. 32 covers a sweep of three image sizes with up to 8 images
+    # per ViT call (18 layouts per rank) with headroom.
+    SGLANG_VIT_CUDA_GRAPH_CACHE_CAPACITY = EnvInt(32)
     # Use the fully-vectorized ViT position-embedding interpolation (no per-image
     # Python loop / CPU<->GPU sync). Bit-exact with the legacy implementation;
     # set False to fall back to the per-image loop.

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

@@ -793,7 +793,7 @@ def forward(
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
-        cu_seqlens: torch.Tensor | SingletonCache | None,
+        cu_seqlens: torch.Tensor | SingletonCache | list | None,
         bsz: int,
         seq_len: int,
         softmax_scale: Optional[float] = None,
@@ -803,6 +803,11 @@ def forward(
         if forward_metadata is not None:
             cu_seqlens_gpu = forward_metadata.cu_seqlens
             max_seqlen = forward_metadata.max_seqlen
+        elif isinstance(cu_seqlens, list):
+            # ViT CUDA graph runners pass [cu_seqlens, max_seqlen]; models without
+            # a runner keep passing tensors even when the graph env var is set.
+            cu_seqlens_gpu = cu_seqlens[0]
+            max_seqlen = cu_seqlens[1]
         else:
             cu_seqlens_gpu = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
             cu_seqlens_gpu = cu_seqlens_gpu.to(dtype=torch.int32).to(q.device)
@@ -1488,9 +1493,11 @@ def forward(
             # CPU-to-CUDA copy. The eager version is captured as part of the
             # graph, so its pointwise work is still replayed without launch
             # overhead.
+            # On ROCm the compiled version captures cleanly after the runner's
+            # eager warmup and matches eager mode bit for bit, so keep it there.
             rotary_fn = (
                 apply_rotary_pos_emb_native_eager
-                if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get()
+                if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get() and not _is_hip
                 else apply_rotary_pos_emb
             )
             q, k = rotary_fn(q, k, cos, sin)

python/sglang/srt/models/qwen3_vl.py

@@ -816,61 +816,6 @@ def bucket_flashinfer_max_seqlen(self, real_max_seqlen: int) -> int:
             round_up(real_max_seqlen, FLASHINFER_MAX_SEQLEN_BUCKETS[-1]),
         )
 
-    def fast_pos_embed_interpolate(self, grid_thw):
-        """Interpolate position embeddings for (batch, 3) size input dimensions.
-
-        Performs bilinear interpolation on spatial dimensions (height, width) and replicates
-        along temporal dimension. The result is reorganized according to spatial_merge_size.
-
-        Args:
-            grid_thw: Tensor of shape [batch_size, 3] with (temporal, height, width) dimensions
-                     in patches for each sample.
-
-        Returns:
-            Interpolated position embeddings tensor.
-        """
-        grid_thw_cpu = grid_thw.cpu().numpy()
-
-        # transfer data to CPU before loop
-        temporal_dims = grid_thw_cpu[:, 0].tolist()
-        height_dims = grid_thw_cpu[:, 1].tolist()
-        width_dims = grid_thw_cpu[:, 2].tolist()
-
-        device = self.pos_embed.weight.device
-        dtype = self.pos_embed.weight.dtype
-
-        patches_size = [h * w for h, w in zip(height_dims, width_dims)]
-        total_patches = sum(patches_size)
-        all_indices_np = np.zeros((4, total_patches), dtype=np.int64)
-        all_weights_np = np.zeros((4, total_patches), dtype=np.float32)
-
-        current_idx = 0
-
-        # calculate indices and weights on CPU
-        for t, h, w in zip(temporal_dims, height_dims, width_dims):
-            h_idxs = self._get_interpolation_indices(h)
-            w_idxs = self._get_interpolation_indices(w)
-
-            indices, weights = self._calculate_indices_and_weights(h_idxs, w_idxs)
-
-            end_idx = current_idx + h * w
-            for i in range(4):
-                all_indices_np[i, current_idx:end_idx] = indices[i]
-                all_weights_np[i, current_idx:end_idx] = weights[i]
-            current_idx = end_idx
-
-        idx_tensor = torch.from_numpy(all_indices_np).to(device)
-        weight_tensor = torch.from_numpy(all_weights_np).to(dtype=dtype, device=device)
-
-        # calculate interpolation
-        pos_embeds = self.pos_embed(idx_tensor.view(-1))
-        pos_embeds = pos_embeds.view(4, total_patches, -1)
-        patch_pos_embeds = (pos_embeds * weight_tensor.unsqueeze(-1)).sum(dim=0)
-        patch_pos_embeds = patch_pos_embeds.split(patches_size)
-        return self._get_position_embedding(
-            patch_pos_embeds, temporal_dims, height_dims, width_dims
-        )
-
     def compute_flashinfer_batch_offsets_packed(
         self,
         token_cu_seqlens: np.ndarray,
@@ -1145,14 +1090,12 @@ def _prepare_graph_inputs(
         else:
             grid_thw_list = grid_thw.tolist()
 
-        if self.align_corners and self._use_vectorized_pos_embed(len(grid_thw_list)):
-            # The vectorized implementation uses linspace coordinates. In graph mode
-            # the legacy fallback honors enable_precise_embedding_interpolation, so
-            # only use the vectorized path when the active graph interpolation mode
-            # is also linspace; otherwise image count would change the output.
+        # Same interpolation as the eager forward (linspace, as the HF reference), so
+        # turning the ViT graph on does not change the image embeddings.
+        if self._use_vectorized_pos_embed(len(grid_thw_list)):
             pos_embeds = self.fast_pos_embed_interpolate_vectorized(grid_thw_list)
         else:
-            pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
+            pos_embeds = self.fast_pos_embed_interpolate_from_list(grid_thw_list)
         x += pos_embeds
 
         # rotary embedding -> (cos, sin)