Add CUDA device_guard to 2nd part of fused attn

turboderp · turboderp · commit e61d4d31d44b · 2023-07-09T14:51:50.000+02:00
diff --git a/exllama_ext/exllama_ext.cpp b/exllama_ext/exllama_ext.cpp
@@ -518,6 +518,7 @@ void q4_attn_2
 {
     TORCH_CHECK_DTYPE(x, kHalf);
     TORCH_CHECK_DTYPE(attn_output, kHalf);
+    const at::cuda::OptionalCUDAGuard device_guard(x.device());
 
     int height = x.size(0);
 

Original file line number	Diff line number	Diff line change
`@@ -518,6 +518,7 @@ void q4_attn_2`
`518`	`518`	`{`
`519`	`519`	`TORCH_CHECK_DTYPE(x, kHalf);`
`520`	`520`	`TORCH_CHECK_DTYPE(attn_output, kHalf);`
	`521`	`+ const at::cuda::OptionalCUDAGuard device_guard(x.device());`
`521`	`522`
`522`	`523`	`int height = x.size(0);`
`523`	`524`