From 587fd15bd272d03f78dcc58a3e771a88d0c40f64 Mon Sep 17 00:00:00 2001 From: Jia Guo Date: Fri, 24 Apr 2026 12:05:16 -0700 Subject: [PATCH] perf: eliminate attention DtoD copy by passing pre-allocated output to FA (#21985) Co-authored-by: Claude Opus 4.6 --- python/sglang/jit_kernel/flash_attention.py | 4 ++++ python/sglang/jit_kernel/flash_attention_v3.py | 4 ++++ .../layers/attention/flashattention_backend.py | 17 +++++++++++++++++ python/sglang/srt/layers/radix_attention.py | 8 +++++++- sgl-kernel/csrc/flash_extension.cc | 4 ++-- sgl-kernel/python/sgl_kernel/flash_attn.py | 6 ++++-- 6 files changed, 38 insertions(+), 5 deletions(-) diff --git a/python/sglang/jit_kernel/flash_attention.py b/python/sglang/jit_kernel/flash_attention.py index 79a9140ae..da6c59dae 100644 --- a/python/sglang/jit_kernel/flash_attention.py +++ b/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 ( diff --git a/python/sglang/jit_kernel/flash_attention_v3.py b/python/sglang/jit_kernel/flash_attention_v3.py index d9b4ba01a..732e789c5 100644 --- a/python/sglang/jit_kernel/flash_attention_v3.py +++ b/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, ) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 7680f2c6a..193ce338c 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index 1386cccbe..a008216a0 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/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 diff --git a/sgl-kernel/csrc/flash_extension.cc b/sgl-kernel/csrc/flash_extension.cc index 376b69283..1fd387f6a 100644 --- a/sgl-kernel/csrc/flash_extension.cc +++ b/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)); diff --git a/sgl-kernel/python/sgl_kernel/flash_attn.py b/sgl-kernel/python/sgl_kernel/flash_attn.py index 05227bad2..e498f1967 100644 --- a/sgl-kernel/python/sgl_kernel/flash_attn.py +++ b/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