[rocm6.3_internal_testing] Fix SWDEV-459623 - 2 (#1629)

Fix faulty conflict merge when cherry-picking aea0386 to rocm6.3_internal_testing --------- Co-authored-by: Xinya Zhang <[email protected]>
ROCm · Oct 10, 2024 · 4df8a01 · 4df8a01
1 parent 7ac294f
commit 4df8a01
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/flash_api.hip b/aten/src/ATen/native/transformers/hip/flash_attn/flash_api.hip
@@ -96,6 +96,10 @@ mha_fwd(const at::Tensor &q,         // batch_size x seqlen_q x num_heads x head
         int window_size_right,
         const bool return_softmax,
         std::optional<at::Generator> gen_) {
+  // Otherwise the kernel will be launched from cuda:0 device
+  // Cast to char to avoid compiler warning about narrowing
+  at::hip::HIPGuardMasqueradingAsCUDA device_guard{(char)q.get_device()};
+
   auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream();
   check_gpu_arch(stream);