[VLM] Refactor ViT CUDA Graph key generation

[yuan-luo]

Motivation

Currently ViT CUDA Graph table key uses seq_len. There's a case that two sequences with identical seq_len, but cu_seqlens are different. cu_seqlens is the segmentation contract for packed sequences: if the contract changes, the graph cannot be reused.

For example we have 2 sequences, each seq has 3 images, the token num of the images are:
seq1: [4, 6, 8]
seq2: [5, 5, 8]
In this case the sequence length are the same, but the cu_seqlens are different.
cu_seqlens1: [0, 4, 10, 18]
cu_seqlens2:[0, 5, 10, 18]
We should use two separate graphs for each sequence.

This PR is to introduce a new graph_key_hint, which is a model-specific hash value related with cu_seqlens and cu_window_seqlens (example for Qwen2.5-VL).

After thinking twice, I reckon it is a general issue for CUDA Graph scenarios such as spec decode, normal decode as well as Piecewise CUDA Graph. The CUDA Graph key uses batch-size, which adopts varlen mechanism in attention backend. They probably encounter the same issue. In this case, the graph buffer has change be ruined, but the case is very rare. So it is extremely difficult to reproduce.

Manually test passed.
Baseline:

root@6996fb46042d:/sgl-workspace/bench_script# bash bench_n_image.sh 
{"id":"58afb10df28d4e5f93af5782a43b1fcb","object":"chat.completion","created":1768134637,"model":"auto","choices":[{"index":0,"message":{"role":"assistant","content":"图中植物是飞廉（学名：Cirsium vulgare），也叫刺儿菜。飞廉是一种常见的多年生草本植物，属于菊科飞廉属。它的叶子边缘有锯齿，茎上有刺，花朵呈紫色或粉红色。飞廉在欧洲和亚洲的一些地区较为常见，但在中国的一些地区也可能生长。","reasoning_content":null,"tool_calls":null},"logprobs":null,"finish_reason":"stop","matched_stop":151645}],"usage":{"prompt_tokens":970,"total_tokens":1047,"completion_tokens":77,"prompt_tokens_details":null,"reasoning_tokens":0},"metadata":{"weight_version":"default"}}
real    0m0.608s
user    0m0.002s
sys     0m0.003s

PR:

root@6996fb46042d:/sgl-workspace/bench_script# bash bench_n_image.sh 
{"id":"8c454ccd6cff49f69330918b57f80820","object":"chat.completion","created":1768135778,"model":"auto","choices":[{"index":0,"message":{"role":"assistant","content":"图中植物是飞廉（学名：Cirsium vulgare），也叫刺儿菜。飞廉是一种常见的野草，广泛分布于欧洲、亚洲和北美洲。它的叶子边缘有锯齿，茎上有刺，花朵呈紫色或粉红色。飞廉虽然外观上可能看起来很美丽，但它的刺可能会对人和动物造成伤害。","reasoning_content":null,"tool_calls":null},"logprobs":null,"finish_reason":"stop","matched_stop":151645}],"usage":{"prompt_tokens":970,"total_tokens":1050,"completion_tokens":80,"prompt_tokens_details":null,"reasoning_tokens":0},"metadata":{"weight_version":"default"}}
real    0m0.727s
user    0m0.004s
sys     0m0.002s

Modifications

Accuracy Tests

Benchmarking and Profiling

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.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Summary of Changes

Hello @yuan-luo, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request improves the robustness and correctness of CUDA Graph utilization for visual language models by refining the graph key generation process. It moves beyond simple sequence length-based keys to incorporate detailed structural information of the input sequences through a new hashing mechanism. This ensures that CUDA graphs are correctly cached and reused only when the underlying computational structure is identical, preventing potential runtime errors and maintaining performance integrity.

Highlights

• Refactored CUDA Graph Key Generation: The mechanism for generating CUDA Graph keys has been enhanced to prevent incorrect graph sharing. Previously, only seq_len was used, which could lead to issues when sequences had the same total length but different internal structures (cu_seqlens).
• Introduced graph_key_hint: A new graph_key_hint parameter has been added, which is a model-specific hash derived from cu_seqlens and cu_window_seqlens. This hint provides a more granular identifier for CUDA graphs, ensuring that graphs are unique for distinct sequence configurations.
• CPU-side Hashing for Efficiency: The computation of cu_seqlens_list and cu_window_seqlens_list and their subsequent hashing are now performed on the CPU. This avoids unnecessary GPU synchronization, maintaining performance while improving correctness.
• TorchDynamo Compatibility: Logic has been added to temporarily disable torch._dynamo during CUDA graph capture. This addresses potential conflicts where Dynamo's tracing queries CUDA RNG state, which is forbidden during graph capture, ensuring stable operation.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request refactors the CUDA graph key generation for Vision Transformer (ViT) models to make it more robust. Previously, the key was based only on seq_len, which could lead to incorrect graph sharing when sequences had the same length but different internal structures (e.g., different cu_seqlens). The new approach introduces a graph_key_hint derived from a hash of cu_seqlens and cu_window_seqlens, ensuring a unique key for each distinct graph structure. The changes are implemented across qwen2_5_vl.py, qwen3_vl.py, and the core vit_cuda_graph_runner.py. The implementation correctly computes the hash on the CPU to avoid performance degradation from GPU synchronization. The code is well-structured and the changes are sound. I have one minor suggestion to remove a redundant operation.

[yuan-luo]

/tag-and-rerun-ci

[yuan-luo]

/rerun-failed-ci

[yuan-luo]

/rerun-failed-ci

[yuan-luo]

/rerun-failed-ci

[yuan-luo]

/rerun-failed-ci