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

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

python/sglang/jit_kernel/flash_attention.py

@@ -41,6 +41,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
@@ -164,6 +165,7 @@ def flash_attn_with_kvcache(
             sm_margin=sm_margin,
             return_softmax_lse=return_softmax_lse,
             sinks=sinks,
+            out=out,
         )
     elif ver == 4:
         from .flash_attention_v4 import (
@@ -232,6 +234,7 @@ def flash_attn_varlen_func(
     score_mod=None,
     aux_tensors=None,
     ver=3,
+    out=None,
 ):
 
     if ver == 3:
@@ -260,6 +263,7 @@ def flash_attn_varlen_func(
             sm_margin=sm_margin,
             return_softmax_lse=return_softmax_lse,
             sinks=sinks,
+            out=out,
         )
     elif ver == 4:
         from .flash_attention_v4 import (

python/sglang/jit_kernel/flash_attention_v3.py

@@ -119,6 +119,7 @@ def flash_attn_with_kvcache(
     sm_margin=0,  # Can be tuned if some SMs are used for communication
     return_softmax_lse=False,
     sinks=None,
+    out=None,
 ):
     if not _is_fa3_supported():
         raise NotImplementedError(
@@ -160,6 +161,7 @@ def flash_attn_with_kvcache(
         sm_margin,
         return_softmax_lse,
         sinks,
+        out=out,
     )
 
 
@@ -189,6 +191,7 @@ def flash_attn_varlen_func(
     sm_margin=0,
     return_softmax_lse=False,
     sinks=None,
+    out=None,
 ):
 
     if not _is_fa3_supported():
@@ -221,4 +224,5 @@ def flash_attn_varlen_func(
         sm_margin,
         return_softmax_lse,
         sinks,
+        out=out,
     )