minor fixing on tuning constraints

hyukn · hyukn · commit 5b5a724c68ed · 2025-11-28T01:47:44.000Z
Signed-off-by: Yukun He &lt;23156053+hyukn@users.noreply.github.com&gt;
diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
@@ -44,6 +44,7 @@ class MoERunner(TunableRunner):
         dynamic_tensor_specs=(DynamicTensorSpec(
             0, 0, get_last_power_of_2_num_tokens_buckets,
             last_positive_power_of_2), ),
+        constraint_specs=(ConstraintSpec(1, 0, lambda shapes: shapes[0][0]), ),
         tune_max_num_tokens=8192,
         inputs_pre_hook=inputs_pre_hook,
     )

Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,7 @@ class MoERunner(TunableRunner):`
`44`	`44`	`dynamic_tensor_specs=(DynamicTensorSpec(`
`45`	`45`	`0, 0, get_last_power_of_2_num_tokens_buckets,`
`46`	`46`	`last_positive_power_of_2), ),`
	`47`	`+ constraint_specs=(ConstraintSpec(1, 0, lambda shapes: shapes[0][0]), ),`
`47`	`48`	`tune_max_num_tokens=8192,`
`48`	`49`	`inputs_pre_hook=inputs_pre_hook,`
`49`	`50`	`)`