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,
     )

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -691,6 +691,12 @@ class FlashAttentionBackend(AttentionBackend):
         if sinks is not None:
             kwargs["sinks"] = sinks
 
+        _fa_out = (
+            forward_batch._attn_output.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+            if getattr(forward_batch, "_attn_output", None) is not None
+            else None
+        )
+
         # Get the appropriate page table based on whether we're using local attention
         if use_local_attn:
             local_metadata = metadata.local_attn_metadata
@@ -797,6 +803,7 @@ class FlashAttentionBackend(AttentionBackend):
                     window_size=window_size,
                     softcap=layer.logit_cap,
                     num_splits=self.num_splits,
+                    out=_fa_out,
                     **kwargs,
                 )
             else:
@@ -817,6 +824,7 @@ class FlashAttentionBackend(AttentionBackend):
                     v_descale=v_descale,
                     return_softmax_lse=use_cascade_attn,
                     num_splits=self.num_splits,
+                    out=_fa_out,
                     ver=self.fa_impl_ver,
                     **kwargs,
                 )
@@ -885,6 +893,7 @@ class FlashAttentionBackend(AttentionBackend):
                         softmax_scale=layer.scaling,
                         causal=False,
                         return_softmax_lse=True,
+                        out=_fa_out,
                         ver=self.fa_impl_ver,
                         **kwargs,
                     )
@@ -911,6 +920,7 @@ class FlashAttentionBackend(AttentionBackend):
                         softmax_scale=layer.scaling,
                         causal=True,
                         return_softmax_lse=forward_batch.mha_return_lse,
+                        out=_fa_out,
                         ver=self.fa_impl_ver,
                         **kwargs,
                     )
@@ -1064,6 +1074,12 @@ class FlashAttentionBackend(AttentionBackend):
         if sinks is not None:
             kwargs["sinks"] = sinks
 
+        _fa_out = (
+            forward_batch._attn_output.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+            if getattr(forward_batch, "_attn_output", None) is not None
+            else None
+        )
+
         k_descale, v_descale = None, None
         # only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention
         # has corresponding quantization method so that layer.k_scale is not None,
@@ -1175,6 +1191,7 @@ class FlashAttentionBackend(AttentionBackend):
                     v_descale=v_descale,
                     return_softmax_lse=use_cascade_attn,
                     num_splits=self.num_splits,
+                    out=_fa_out,
                     ver=self.fa_impl_ver,
                     scheduler_metadata=sched_meta,
                     **kwargs,

python/sglang/srt/layers/radix_attention.py

@@ -190,6 +190,11 @@ def unified_attention_with_output(
         if hasattr(token_to_kv_pool, "set_swa_loc"):
             token_to_kv_pool.set_swa_loc(forward_batch.out_cache_loc_swa)
 
+    # Store pre-allocated output for FA backend to write directly into.
+    # Must slice to real_num_tokens to match the narrowed query shape —
+    # the FA kernel validates out.size(0) == q.size(0).
+    forward_batch._attn_output = output[:real_num_tokens]
+
     ret = forward_batch.attn_backend.forward(
         query,
         key,
@@ -206,7 +211,8 @@ def unified_attention_with_output(
     ):
         token_to_kv_pool.set_swa_loc(original_swa_loc)
 
-    output[:real_num_tokens].view(ret.shape).copy_(ret)
+    if ret.data_ptr() != output.data_ptr():
+        output[:real_num_tokens].view(ret.shape).copy_(ret)
     return