[GENERAL] add device index in vllmGetQueue (vllm-project#61)

Signed-off-by: Yizhou Wang <yizhou.wang@intel.com>

Author: YizhouZ

csrc/flash_attn/flash_api.cpp

@@ -21,7 +21,7 @@ std::vector<at::Tensor> mha_varlen_fwd(
     bool is_causal, int window_size_left, int window_size_right,
     const float softcap, const bool return_softmax,
     std::optional<at::Generator> gen_) {
-  auto& queue = vllm::xpu::vllmGetQueue();
+  auto& queue = vllm::xpu::vllmGetQueue(q.device().index());
 
   at::Tensor out;
   if (out_.has_value()) {

csrc/utils.h

@@ -7,8 +7,8 @@
 namespace vllm {
 namespace xpu {
 
-static inline sycl::queue& vllmGetQueue() {
-  auto current_stream = c10::xpu::getCurrentXPUStream();
+static inline sycl::queue& vllmGetQueue(at::DeviceIndex device_index = -1) {
+  auto current_stream = c10::xpu::getCurrentXPUStream(device_index);
   auto& queue = current_stream.queue();
   return queue;
 }