File tree Expand file tree Collapse file tree 1 file changed +4
-2
lines changed Expand file tree Collapse file tree 1 file changed +4
-2
lines changed Original file line number Diff line number Diff line change @@ -2101,7 +2101,9 @@ def check_torchair_graph_batch_sizes(self):
2101
2101
if self .parallel_config .enable_expert_parallel :
2102
2102
new_graph_batch_sizes = []
2103
2103
for graph_batch_size in self .torchair_graph_batch_sizes :
2104
- cur_graph_batch_size = graph_batch_size + tp_size - graph_batch_size % tp_size
2105
- if cur_graph_batch_size not in new_graph_batch_sizes :
2104
+ cur_graph_batch_size = (graph_batch_size + tp_size -
2105
+ 1 ) // tp_size * tp_size
2106
+ if cur_graph_batch_size not in new_graph_batch_sizes and \
2107
+ cur_graph_batch_size <= self .scheduler_config .max_num_batched_tokens :
2106
2108
new_graph_batch_sizes .append (cur_graph_batch_size )
2107
2109
self .torchair_graph_batch_sizes = new_graph_batch_sizes
You can’t perform that action at this time.
0 commit comments