[chore] Update test_beam_search to correctly use padding, when necessary

stnie · stnie · commit a92b1bd4d2a9 · 2025-08-06T10:24:24.000Z
Signed-off-by: Stefan Niebler &lt;82932102+stnie@users.noreply.github.com&gt;
diff --git a/tests/unittest/_torch/test_beam_search.py b/tests/unittest/_torch/test_beam_search.py
@@ -63,7 +63,8 @@ def llm_cuda_graph(fixed_params, input_prompts):
         enable_trtllm_sampler=True,
         max_beam_width=fixed_params["max_beam_width"],
         disable_overlap_scheduler=False,
-        cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2, 4, 8]),
+        cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2, 4, 8],
+                                          enable_padding=True),
     )
 
 

Original file line number	Diff line number	Diff line change
`@@ -63,7 +63,8 @@ def llm_cuda_graph(fixed_params, input_prompts):`
`63`	`63`	`enable_trtllm_sampler=True,`
`64`	`64`	`max_beam_width=fixed_params["max_beam_width"],`
`65`	`65`	`disable_overlap_scheduler=False,`
`66`		`- cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2, 4, 8]),`
	`66`	`+ cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2, 4, 8],`
	`67`	`+ enable_padding=True),`
`67`	`68`	`)`
`68`	`69`
`69`	`70`