minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Jia Guo
2026-04-24 12:05:16 -07:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 6d03861476
commit 587fd15bd2
6 changed files with 38 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/flash_attention.py
+4
View File
@@ -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
+4
View File
@@ -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
+17
View File
@@ -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
+7 -1
View File
@@ -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
sgl-kernel/csrc/flash_extension.cc
+2 -2
View File
@@ -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
+4 -2
View File
@@ -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