[CI] Stabilize batch submission in full CUDA graph tests (#58810)
Andreas Karatzas committed
Sep 26, 2026 at 10:35 UTC
0239bd2549224a4bc4245ba5a70e7acb120cbeb7
1 file changed
+13
-2
tests/compile/fullgraph/test_full_cudagraph.py
+13
-2
@@ -10,6 +10,7 @@ from tests.utils import wait_for_gpu_memory_to_clear
10
from tests.v1.attention.utils import full_cg_backend_configs as backend_configs
11
from vllm import LLM, SamplingParams
12
from vllm.config import CompilationConfig, CUDAGraphMode
13
+from vllm.outputs import RequestOutput
14
from vllm.platforms import current_platform
15
from vllm.utils.torch_utils import is_torch_equal_or_newer
16
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -169,10 +170,20 @@ class TestFullCUDAGraph:
170
temperature=0.0, max_tokens=max_tokens, top_p=1.0
171
)
172
172
- piecewise_responses = piecewise_llm.generate(prompts, sampling_params)
173
- full_responses = full_cudagraph_llm.generate(prompts, sampling_params)
173
+ responses = []
174
+ for llm in (piecewise_llm, full_cudagraph_llm):
175
+ # Queue the whole batch before scheduling so both engines exercise
176
+ # the requested batch size instead of timing-dependent sub-batches.
177
+ llm.sleep(level=0, mode="keep")
178
+ try:
179
+ llm.enqueue(prompts, sampling_params)
180
+ finally:
181
+ llm.wake_up(tags=["scheduling"])
182
+ responses.append(llm.wait_for_completion(output_type=RequestOutput))
183
+ piecewise_responses, full_responses = responses
184
185
# Check that all responses are the same
186
+ assert len(piecewise_responses) == len(full_responses) == batch_size
187
for piecewise_res, full_res in zip(piecewise_responses, full_responses):
188
assert (
189
piecewise_res.outputs[0].text.lower()