In the OSDI paper, the Pipeline Structure Search process is described as:
In the first stage, auto-search takes as input the dense batch size, the operation dependencies, and the profiling of interference-free kernels, and then produces as output, the number, batch size, and order of each nano-operation using mixed integer linear programming (MILP).
and
Thus, auto-search begins its search by dividing all operations into two nano-operations and formulating the MILP problem to identify the batch size, execution time, and ordering of the nano-operations. If the temporary best solution has bubbles for compute, auto-search increases the number of nano-operations for operations near the bubble to improve resource utilization until MILP cannot produce better solutions.
Based on my understanding, the input specifies how many nano-operations should every operator be split into, while the specific batch size of each split nano-operation should be calculated by the MILP model.
However, it seems like the MILP model in new_search.py uses a fixed batch size for each operation(specifically, decode and prefill batch size for two nano-operations) instead of treating the batch size as an variable, assuming a pre-defined batch size input for each nano-operation.
I'm curious if this simplified modeling is enough, and how should I decide the batch size input for each nano-operation in this case.
In the OSDI paper, the Pipeline Structure Search process is described as:
In the first stage, auto-search takes as input the dense batch size, the operation dependencies, and the profiling of interference-free kernels, and then produces as output, the number, batch size, and order of each nano-operation using mixed integer linear programming (MILP).
and
Thus, auto-search begins its search by dividing all operations into two nano-operations and formulating the MILP problem to identify the batch size, execution time, and ordering of the nano-operations. If the temporary best solution has bubbles for compute, auto-search increases the number of nano-operations for operations near the bubble to improve resource utilization until MILP cannot produce better solutions.
Based on my understanding, the input specifies how many nano-operations should every operator be split into, while the specific batch size of each split nano-operation should be calculated by the MILP model.
However, it seems like the MILP model in
new_search.pyuses a fixed batch size for each operation(specifically, decode and prefill batch size for two nano-operations) instead of treating the batch size as an variable, assuming a pre-defined batch size input for each nano-operation.I'm curious if this simplified modeling is enough, and how should I decide the batch size input for each nano-operation in this case.