perf: eliminate attention DtoD copy by passing pre-allocated output to FA (#21985)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

sgl-kernel/csrc/flash_extension.cc

@@ -29,7 +29,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "    Tensor?  k_new,"             // (b, s_k_new, h_k, d) or (total_k_new, h_k, d)
       "    Tensor?  v_new,"             // (b, s_k_new, h_k, dv) or (total_k_new, h_k, dv)
       "    Tensor?  q_v,"               // (b, s_q, h, dv) or (total_q_new, h, dv)
-      "    Tensor?  out,"               // (b, s_q, h, dv) or (total_q, h, dv)
+      "    Tensor(a!)?  out,"           // (b, s_q, h, dv) or (total_q, h, dv)
       "    Tensor?  cu_seqlens_q,"      // b+1
       "    Tensor?  cu_seqlens_k,"      // b+1
       "    Tensor?  cu_seqlens_k_new,"  // b+1
@@ -58,7 +58,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "    bool?    pack_gqa,"
       "    int      sm_margin,"
       "    Tensor?  sinks"
-      ") -> (Tensor, Tensor, Tensor, Tensor)");  // NEW return type: tuple of 4 tensors
+      ") -> (Tensor(a!), Tensor, Tensor, Tensor)");  // first return aliases out
 
   m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));
 

sgl-kernel/python/sgl_kernel/flash_attn.py

@@ -68,6 +68,7 @@ def flash_attn_with_kvcache(
     score_mod=None,
     aux_tensors=None,
     ver=3,
+    out=None,
 ):
     """
     If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
@@ -193,7 +194,7 @@ def flash_attn_with_kvcache(
         k,
         v,
         qv,
-        None,  # out
+        out,  # out (pre-allocated output to avoid DtoD copy)
         cu_seqlens_q,
         None,  # cu_seqlens_k
         cu_seqlens_k_new,
@@ -256,6 +257,7 @@ def flash_attn_varlen_func(
     score_mod=None,
     aux_tensors=None,
     ver=3,
+    out=None,
 ):
 
     if not is_fa3_supported():
@@ -280,7 +282,7 @@ def flash_attn_varlen_func(
         None,  # k_new
         None,  # v_new
         qv,  # qv
-        None,  # out
+        out,  # out
         cu_seqlens_q,
         cu_seqlens_k,
         None,  # cu_seqlens_k_new