[vorapolsiloai]

Motivation

On ROCm with SGLANG_USE_AITER=1, the vision encoder resolves to the aiter_attn backend. ViTCudaGraphRunner only knew the triton_attn, fa3 and fa4 argument layouts, and VisionAiterAttention could not take the runner's list argument. With SGLANG_VIT_ENABLE_CUDA_GRAPH=1, the first graph warmup/capture therefore failed. As a result, ViT CUDA graphs could not be used on AMD Instinct GPUs. The target workload is Qwen3-VL-family models (Qwen3.8-27B) on MI355X, where image requests are short-output and spend much of their time in the encoder.

Once the graph ran, its output was not identical to eager mode, it could fail to capture while the server was busy, and its memory grew with every new image layout. The follow-up commits fix those so the graph is a drop-in for eager.

Modifications

Rebased onto current main.

aiter_attn support (first commit)

• ViTCudaGraphRunner: warm up and capture aiter_attn with the [cu_seqlens, max_len] layout already used for FA3/FA4.
• VisionAiterAttention.forward: when the runner passes [cu_seqlens, max_seqlen], use them directly, as VisionFlash4Attention does. Graph replay then does no .item() host sync. Tensor callers, including models without a graph runner while the env var is set, keep the existing eager path unchanged.

Graph output == eager output

• Qwen3-VL position embeddings: forward_with_cuda_graph interpolated with fast_pos_embed_interpolate, which follows enable_precise_embedding_interpolation (half-pixel coordinates by default), while the eager forward always uses the linspace interpolation. The graph path now uses the same interpolation as eager, and the Qwen3-VL fast_pos_embed_interpolate it called, now unused, is removed.
• Rotary workspace dtype: the runner copied the model's cos/sin into workspaces of the ViT dtype, so fp32 cos/sin (Qwen3-VL) were rounded to bf16 in graph mode only. The workspace now keeps the dtype the model passes.
• Rotary function on ROCm: in graph mode VisionAttention used the eager rotary helper while eager mode uses the torch.compiled one, so outputs differed in the last bits. On ROCm the compiled helper captures cleanly after the runner's eager block-0 warmup, so it is used in both modes there. Other platforms keep the eager helper from fix: fix vlm cuda graph shape stability #30868 (compiling inside the capture failed there); this PR does not change their behavior.

Robustness and memory

• Thread-local capture: ViT graphs are captured on the first request of each layout, while the process is serving. Other threads (e.g. the c10d NCCL watchdog) keep calling the runtime during capture, which the default global capture mode turns into a capture error. Capture now uses capture_error_mode="thread_local".
• Shared graph pool: each graph had its own private memory pool, so reserved memory grew with every new image layout. All ViT graphs now share one pool, as the LLM graph runners and the Kimi-K3 ViT runner do.
• Bounded graph cache: graphs are keyed on the exact attention layout, so new image sizes and batch compositions add graphs without limit. New SGLANG_VIT_CUDA_GRAPH_CACHE_CAPACITY (default 32) caps the number of captured graphs; later layouts run the same blocks + merger code eagerly (shared _forward_blocks helper, no per-layout buffers), and captured layouts keep replaying. Sizing: a sweep of three image sizes (512², 768×1024, 2048×1600) with up to 8 images per ViT call needs 18 graphs per rank; each graph adds only its input, output and attention-workspace buffers (about 7 KB per patch for a 1152-wide encoder, about 180 MB for the largest layout), while activations share one pool. The model eager path (SGLANG_VIT_ENABLE_CUDA_GRAPH unset) is unchanged.

Tests and docs

• Unit tests (CPU): aiter list/tensor inputs with the graph env var on and off; rotary workspace keeps fp32 cos/sin for a bf16 ViT (fails before the fix); cache capacity runs new layouts uncaptured without allocating per-layout buffers, still replays captured layouts, and is unbounded by default.
• Docs: aiter_attn in the backend layout list and the cache capacity in cuda_graph_for_multi_modal_encoder.mdx; the new env var in environment_variables.mdx.

ViT CUDA graphs stay opt-in (SGLANG_VIT_ENABLE_CUDA_GRAPH, default off).

Dependency: #43315 (still open). Graph replay returns the graph's static output buffer, which a later replay can overwrite. #43315 (return an owned output from ViT graph replays) fixes this for all backends. This PR does not duplicate it. The shared graph pool depends on it more strongly: with one pool, a replay of a different graph can also reuse that memory, so the pool commit is only safe with #43315 applied. All validation of this PR is done with #43315 applied, and it should merge after #43315.

Accuracy Tests

On an earlier internal MI355X build with older main + #43315 + this PR's changes except the cache bound, the ViT graph output was bit-exact vs eager in 148 of 148 checks on Qwen3.8-27B.

To be re-run on this head (MI355X): ViT graph vs eager output for several image sizes and batch compositions, plus MMMU with the graph off/on.

Speed Tests and Profiling

On that earlier build, enabling the ViT graph raised image-serving throughput on MI355X full sweeps (Qwen3.8-27B) by about +8% at TP1 and +10.5% at TP2. The TP2 build also carried a separate batched-ViT change that will be its own PR.

MI355X numbers to be re-measured on this head: image serving rows (bench_serve, random images 512² / 768×1024 / 2048×1600, OSL 64) with the graph off/on, plus capture count and memory at high concurrency with and without SGLANG_VIT_CUDA_GRAPH_CACHE_CAPACITY.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #40418501314
Latest PR Test (Extra): ❌ Run #40418501286
Latest PR Test (AMD ROCm 10): ❌ Run #40418501264

[mintlify]

Preview approval required

This pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit 882c41f.

Approve preview