Allow mixed-batch sampling in dynamic inference - #1927
Conversation
tdene
added
the
Expert Review
santhnm2
left a comment
santhnm2
left a comment
There was a problem hiding this comment.
LGTM, just left a few comments about keeping arguments as optional for now to preserve backwards compatibility.
| new_sample = torch.zeros(batch_size, dtype=torch.int64, device=last_token_logits.device) | ||
| termination_id = torch.zeros_like(new_sample, dtype=torch.int64) | ||
|
|
||
| for sampling_params, mask in active_sampling_map: |
There was a problem hiding this comment.
Can we avoid this loop when not necessary? in cases where users dont have sampling params per request ? (maybe add a paramter that users can set which notes that the same sampling params will be used for all requests ? )
Will there be a performance difference ?
There was a problem hiding this comment.
If there's only 1 sets of sampling params, the active_sampling_map will only have 1 entry. We'll only go through the loop once in this particular piece of code, so there should be no performance difference here.
But you are right, however, that there is a loop that cannot be avoided in dynamic_engine.py L598, here. That loop only runs once per step and loops through every currently-active request.
I have a later PR in this series that solves this performance issues by overlapping the CPU and GPU compute. But I will also solve the issue in this PR, instead of leaving it for later. Let me think and address this.
There was a problem hiding this comment.
Addressed in the most recent commit. Does that seem good @shanmugamr1992, as a temporary solution until future PRs optimize this even more?
6 participants
This MR allows for dynamic inference to process requests with non-uniform sampling parameters.
Future MRs will provide large optimizations.