A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 18:45 UTC
[MM][CG] Support `--enable-vit-cuda-graph` option for VLM examples (#40580)
Signed-off-by: shen-shanshan <467638484@qq.com>
Shanshan Shen committed
Apr 23, 2026 at 13:46 UTC
fe57be7809672e5c4d100b55ce8649dd34d3bbc0
2 files changed
+40
-6
examples/offline_inference/vision_language.py
+35
-5
@@ -2463,6 +2463,12 @@ MODELS_NEED_VIDEO_METADATA = [
2463
]
2464
2465
2466
+MODELS_SUPPORT_VIT_CUDA_GRAPH = [
2467
+ "qwen3_vl",
2468
+ "qwen3_vl_moe",
2469
+]
2470
+
2471
+
2472
def get_multi_modal_input(args):
2473
"""
2474
return {
@@ -2575,6 +2581,29 @@ def apply_image_repeat(
2581
return inputs, inputs_with_empty_media
2582
2583
2584
+def maybe_add_vit_cuda_graph_compilation_config(args, engine_args):
2585
+ model = args.model_type
2586
+ modality = args.modality
2587
+ enable_vit_cuda_graph = args.enable_vit_cuda_graph
2588
+
2589
+ if enable_vit_cuda_graph and model in MODELS_SUPPORT_VIT_CUDA_GRAPH:
2590
+ if modality == "image" or modality == "video":
2591
+ vision_items_per_batch = 1
2592
+ elif modality == "image+video":
2593
+ vision_items_per_batch = 2
2594
+ else:
2595
+ raise ValueError(
2596
+ f"modality={modality} is not supported for vit cuda graph."
2597
+ )
2598
+
2599
+ engine_args.compilation_config = {
2600
+ "cudagraph_mm_encoder": True,
2601
+ "encoder_cudagraph_max_vision_items_per_batch": vision_items_per_batch,
2602
+ }
2603
+
2604
+ return engine_args
2605
+
2606
+
2607
@contextmanager
2608
def time_counter(enable: bool):
2609
if enable:
@@ -2625,33 +2654,28 @@ def parse_args():
2654
default=0,
2655
help="Set the seed when initializing `vllm.LLM`.",
2656
)
2628
-
2657
parser.add_argument(
2658
"--image-repeat-prob",
2659
type=float,
2660
default=None,
2661
help="Simulates the hit-ratio for multi-modal preprocessor cache (if enabled)",
2662
)
2635
-
2663
parser.add_argument(
2664
"--disable-mm-processor-cache",
2665
action="store_true",
2666
help="If True, disables caching of multi-modal processor.",
2667
)
2641
-
2668
parser.add_argument(
2669
"--time-generate",
2670
action="store_true",
2671
help="If True, then print the total generate() call time",
2672
)
2647
-
2673
parser.add_argument(
2674
"--use-different-prompt-per-request",
2675
action="store_true",
2676
help="If True, then use different prompt (with the same multi-modal "
2677
"data) for each request.",
2678
)
2654
-
2679
parser.add_argument(
2680
"--verify-mm-cache-hit-with-uuids",
2681
action="store_true",
@@ -2665,6 +2689,11 @@ def parse_args():
2689
default=None,
2690
help="Tensor parallel size to override the model's default setting. ",
2691
)
2692
+ parser.add_argument(
2693
+ "--enable-vit-cuda-graph",
2694
+ action="store_true",
2695
+ help="If True, will enable vit cuda graph capture and replay for the model.",
2696
+ )
2697
return parser.parse_args()
2698
2699
@@ -2698,6 +2727,7 @@ def main(args):
2727
engine_args.mm_processor_cache_gb = mm_processor_cache_gb
2728
if args.tensor_parallel_size is not None:
2729
engine_args.tensor_parallel_size = args.tensor_parallel_size
2730
+ engine_args = maybe_add_vit_cuda_graph_compilation_config(args, engine_args)
2731
llm = LLM.from_engine_args(engine_args)
2732
2733
# Don't want to check the flag multiple times, so just hijack `prompts`.
vllm/model_executor/models/qwen3_vl.py
+5
-1
@@ -1802,7 +1802,11 @@ class Qwen3VLForConditionalGeneration(
1802
# spatial_merge_size=2 → 8x8 = 64 tokens
1803
min_budget = 64
1804
# Max: capped by max_num_batched_tokens
1805
- max_budget = vllm_config.scheduler_config.max_num_batched_tokens
1805
+ # TODO(shen-shanshan): the max_budget auto-infer needs to be optimized later.
1806
+ max_budget = min(
1807
+ vllm_config.scheduler_config.max_num_batched_tokens,
1808
+ self.model_config.max_model_len,
1809
+ )
1810
return (min_budget, max_budget)
1811
1812
def _get_pixel_values_by_modality(