resolve comments

Signed-off-by: Ceng23333 <441651826@qq.com>

Author: Ceng23333

src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc

@@ -45,24 +45,24 @@ void *plan(Tensor out,
 
 void run(void *planned_meta) {
 #ifdef ENABLE_FLASH_ATTN
-#ifdef ENABLE_NVIDIA_API
-    c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream());
-#elif defined(ENABLE_METAX_API)
+#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API)
     c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream());
 #endif
     auto *p = reinterpret_cast<PlannedMeta *>(planned_meta);
 
     // Paged KV caches must be contiguous for flash-attn; avoid extra copies for q/metadata when already dense.
-    auto out_at = infinicore::adaptor::to_aten_tensor(p->out);
-    const bool out_need_copy_back = !out_at.is_contiguous();
-    auto out_tensor = out_need_copy_back ? out_at.contiguous() : out_at;
+    const bool out_need_copy_back = !p->out->is_contiguous();
+    Tensor out_work = out_need_copy_back ? p->out->contiguous() : Tensor(p->out);
+    auto out_tensor = infinicore::adaptor::to_aten_tensor(out_work);
     auto q = infinicore::adaptor::to_aten_tensor(p->q);
 #if defined(ENABLE_NVIDIA_API)
     auto k_cache = infinicore::adaptor::to_aten_tensor(p->k_cache);
     auto v_cache = infinicore::adaptor::to_aten_tensor(p->v_cache);
 #elif defined(ENABLE_QY_API) || defined(ENABLE_METAX_API)
-    auto k_cache = infinicore::adaptor::to_aten_tensor(p->k_cache).contiguous();
-    auto v_cache = infinicore::adaptor::to_aten_tensor(p->v_cache).contiguous();
+    Tensor k_cache_work = p->k_cache->contiguous();
+    Tensor v_cache_work = p->v_cache->contiguous();
+    auto k_cache = infinicore::adaptor::to_aten_tensor(k_cache_work);
+    auto v_cache = infinicore::adaptor::to_aten_tensor(v_cache_work);
 #endif
     auto seqlens_k = std::optional<const at::Tensor>(infinicore::adaptor::to_aten_tensor(p->seqlens_k));
     auto block_table = std::optional<at::Tensor>(infinicore::adaptor::to_aten_tensor(p->block_table));
@@ -119,7 +119,7 @@ void run(void *planned_meta) {
         out_tensor.copy_(result[0]);
     }
     if (out_need_copy_back) {
-        out_at.copy_(out_tensor);
+        p->out->copy_from(out_work);
     }
 #else
     throw std::runtime_error("FlashAttention is not enabled in this build");

src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc

@@ -49,12 +49,13 @@ namespace {
 
 #ifdef ENABLE_FLASH_ATTN
 struct VarlenFlashPrepared {
+    Tensor k_work;
+    Tensor v_work;
+    Tensor out_work_ic;
     at::Tensor q;
     at::Tensor k;
     at::Tensor v;
-    at::Tensor out_at;
     bool out_need_copy_back;
-    at::Tensor out_work;
     std::optional<at::Tensor> out_opt;
     at::Tensor cu_seqlens_q;
     at::Tensor cu_seqlens_kv;
@@ -69,12 +70,15 @@ VarlenFlashPrepared prepare_varlen_flash_tensors(PlannedMeta *p) {
     VarlenFlashPrepared t;
     // Varlen flash-attn: keep k/v contiguous for dense/paged layout; avoid extra copies for q/metadata when already dense.
     t.q = infinicore::adaptor::to_aten_tensor(p->q);
-    t.k = infinicore::adaptor::to_aten_tensor(p->k).contiguous();
-    t.v = infinicore::adaptor::to_aten_tensor(p->v).contiguous();
-    t.out_at = infinicore::adaptor::to_aten_tensor(p->out);
-    t.out_need_copy_back = !t.out_at.is_contiguous();
-    t.out_work = t.out_need_copy_back ? t.out_at.contiguous() : t.out_at;
-    t.out_opt = std::optional<at::Tensor>(t.out_work);
+    t.k_work = p->k->contiguous();
+    t.v_work = p->v->contiguous();
+    t.k = infinicore::adaptor::to_aten_tensor(t.k_work);
+    t.v = infinicore::adaptor::to_aten_tensor(t.v_work);
+
+    t.out_need_copy_back = !p->out->is_contiguous();
+    t.out_work_ic = t.out_need_copy_back ? p->out->contiguous() : Tensor(p->out);
+    auto out_work = infinicore::adaptor::to_aten_tensor(t.out_work_ic);
+    t.out_opt = std::optional<at::Tensor>(out_work);
     t.cu_seqlens_q = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_q);
     t.cu_seqlens_kv = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_k);
     t.block_table = std::optional<at::Tensor>(infinicore::adaptor::to_aten_tensor(p->block_table));
@@ -87,9 +91,9 @@ VarlenFlashPrepared prepare_varlen_flash_tensors(PlannedMeta *p) {
     return t;
 }
 
-void copy_varlen_flash_output_back(VarlenFlashPrepared &t) {
+void copy_varlen_flash_output_back(PlannedMeta *p, VarlenFlashPrepared &t) {
     if (t.out_need_copy_back) {
-        t.out_at.copy_(t.out_work);
+        p->out->copy_from(t.out_work_ic);
     }
 }
 
@@ -135,7 +139,7 @@ void run_flashattn_varlen_metax(PlannedMeta *p) {
         flash_attn_mars_ext
 #endif
     );
-    copy_varlen_flash_output_back(t);
+    copy_varlen_flash_output_back(p, t);
 }
 #endif