[Misc] Various simplifications and typing fixes (#5368)

2024-06-10 19:29:02 -07:00
parent 76477a93b7
commit a008629807
8 changed files with 62 additions and 89 deletions
--- a/vllm/model_executor/layers/rejection_sampler.py
+++ b/vllm/model_executor/layers/rejection_sampler.py
@@ -306,8 +306,10 @@ class RejectionSampler(nn.Module):

        # Fill in the first k columns of the output tensor using masks and data
        # tensors.
-        output[:, :k] = torch.where(accepted_mask, draft_token_ids,
-                                    -torch.ones_like(draft_token_ids))
+        torch.where(accepted_mask,
+                    draft_token_ids,
+                    -torch.ones_like(draft_token_ids),
+                    out=output)

        # Fill the last column.
        # We check output directly as accepted may have True values inconsistent