[CI] Reduce CUDA graph mode test overhead (#58749)
Misha Goin committed
Sep 26, 2026 at 11:22 UTC
79ac28c5bd363d243ff2dea39306b568f59d4153
1 file changed
+10
-6
tests/v1/cudagraph/test_cudagraph_mode.py
+10
-6
@@ -54,13 +54,16 @@ def test_backend_and_cudagraph_mode_combo(backend_name, cudagraph_mode, supporte
54
55
llm = LLM(
56
model="Qwen/Qwen2-1.5B-Instruct",
57
- max_num_seqs=256,
57
+ max_num_seqs=16,
58
+ max_num_batched_tokens=128,
59
trust_remote_code=True,
60
gpu_memory_utilization=0.45,
61
max_model_len=1024,
62
attention_config=attention_config,
63
compilation_config=CompilationConfig(
63
- mode=CompilationMode.VLLM_COMPILE, cudagraph_mode=cudagraph_mode
64
+ mode=CompilationMode.VLLM_COMPILE,
65
+ cudagraph_mode=cudagraph_mode,
66
+ cudagraph_capture_sizes=[1, 2, 4, 8, 16, 64, 128],
67
),
68
)
69
llm.generate(["Hello, my name is"] * 10)
@@ -72,11 +75,9 @@ attn_backend = "RocmAttn" if current_platform.is_rocm() else "FA2"
75
76
combo_cases_2 = [
77
(attn_backend, "FULL", CompilationMode.NONE, True),
75
- (attn_backend, "FULL", CompilationMode.VLLM_COMPILE, True),
78
(attn_backend, "PIECEWISE", CompilationMode.NONE, True),
79
(attn_backend, "PIECEWISE", CompilationMode.VLLM_COMPILE, True),
80
(attn_backend, "FULL_AND_PIECEWISE", CompilationMode.NONE, True),
79
- (attn_backend, "FULL_AND_PIECEWISE", CompilationMode.VLLM_COMPILE, True),
81
(attn_backend, "FULL_DECODE_ONLY", CompilationMode.NONE, True),
82
(attn_backend, "FULL_DECODE_ONLY", CompilationMode.VLLM_COMPILE, True),
83
(attn_backend, "NONE", CompilationMode.NONE, True),
@@ -100,13 +101,16 @@ def test_cudagraph_compilation_combo(
101
102
llm = LLM(
103
model="Qwen/Qwen2-1.5B-Instruct",
103
- max_num_seqs=256,
104
+ max_num_seqs=16,
105
+ max_num_batched_tokens=128,
106
trust_remote_code=True,
107
gpu_memory_utilization=0.45,
108
max_model_len=1024,
109
attention_config=attention_config,
110
compilation_config=CompilationConfig(
109
- mode=compilation_mode, cudagraph_mode=cudagraph_mode
111
+ mode=compilation_mode,
112
+ cudagraph_mode=cudagraph_mode,
113
+ cudagraph_capture_sizes=[1, 2, 4, 8, 16, 64, 128],
114
),
115
)
116
llm.generate(["Hello, my name is"] * 10)