Ensure logit padding happens on default stream

Author: turboderp

exllamav2/model.py

@@ -989,6 +989,9 @@ def forward_chunk(self,
         if self.tp_context:
             self.tp_context.wait_streams()
 
+        if x is not None and x.is_cuda:
+            torch.cuda.set_stream(torch.cuda.default_stream(x.device))
+
         # Apply logit scale
 
         # if x is not None and self.config.logit_scale != 1: