diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 2ee8d29ccca..d86c19a01b5 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -1656,9 +1656,11 @@ def _( class AllReduceRunner(TunableRunner): tuning_config = TuningConfig( - dynamic_tensor_specs=(DynamicTensorSpec( - 0, 0, get_last_power_of_2_num_tokens_buckets(8192), - last_positive_power_of_2), ), + dynamic_tensor_specs=( + # TODO: change to dynamic shape triggering + DynamicTensorSpec(0, 0, + get_last_power_of_2_num_tokens_buckets(8192), + last_positive_power_of_2), ), constraint_specs=(ConstraintSpec(1, 0, lambda shapes: shapes[0][0]), ), distributed_tuning_strategy=DistributedTuningStrategy.MERGE, )