Skip to content

Commit 5e272ee

Browse files
authored
feat : reduce trt engine build time in testing (NVIDIA#3014)
Signed-off-by: peaceh <[email protected]>
1 parent 1ac0566 commit 5e272ee

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

tests/unittest/attention/test_gpt_attention.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -754,7 +754,9 @@ def _construct_execution(
754754
precision=dtype,
755755
int8=int8_trt_flag,
756756
quant_mode=quant_mode)
757-
757+
# Reuce the TRT engine build time by setting the max allowed number of tactics in builder tactic profiling.
758+
if builder_config.trt_builder_config.max_num_tactics == -1:
759+
builder_config.trt_builder_config.max_num_tactics = 30
758760
if session is None:
759761
engine = builder.build_engine(net, builder_config)
760762
session = tensorrt_llm.runtime.Session.from_serialized_engine(

0 commit comments

Comments
 (0)