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

[fa] Make the FlashAttention backend extensible by subclasses (#33426)

Browse Source 
Signed-off-by: Kurt Shuster <kurt@thinkingmachines.ai>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
Kurt Shuster
2026-09-13 21:21:25 -07:00
committed by GitHub
co-authored by Baizhou Zhang
parent 39e147443b
commit f2111715cd
2 changed files with 132 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/flashattention_backend.py
+32 -11
View File
@@ -50,11 +50,6 @@ if TYPE_CHECKING:
from sgl_kernel import merge_state_v2
from sglang.kernels.ops.attention.flash_attention import (
flash_attn_varlen_func,
flash_attn_with_kvcache,
)
def _should_disable_scheduler_metadata_precompute() -> bool:
return bool(get_parallel().enable_prefill_cp or get_parallel().enable_dp_attention)
@@ -1283,7 +1278,13 @@ class FlashAttentionBackend(AttentionBackend):
aux_tensors=None,
rel_bias=None,
rel_bias_event=None,
):
# Returns (output, lse) with lse in [total_q, num_heads].
return_lse: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
lse_out = None
# Bound in __init__ so a subclass can substitute a different FA4 build.
flash_attn_with_kvcache = self.flash_attn_with_kvcache
flash_attn_varlen_func = self.flash_attn_varlen_func
if score_mod is not None and self.fa_impl_ver != 4:
raise RuntimeError("score_mod is only supported by the FA4 backend.")
cp_active = is_cp_active(forward_batch)
@@ -1563,7 +1564,7 @@ class FlashAttentionBackend(AttentionBackend):
causal=False if use_cascade_attn else causal,
window_size=window_size,
softcap=layer.logit_cap,
return_softmax_lse=use_cascade_attn,
return_softmax_lse=use_cascade_attn or return_lse,
num_splits=self.num_splits,
out=_fa_out,
ver=self.fa_impl_ver,
@@ -1623,6 +1624,8 @@ class FlashAttentionBackend(AttentionBackend):
o_expand,
softmax_lse_expand.T.contiguous(),
)
elif return_lse:
o, lse_out, *_ = result
else:
o = result
else:
@@ -1823,7 +1826,12 @@ class FlashAttentionBackend(AttentionBackend):
else:
o = result
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
o = o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
if return_lse:
assert lse_out is not None
# The varlen kernel emits LSE head-major [num_heads, total_q].
return o, lse_out.transpose(0, 1).contiguous()
return o
def forward_decode(
self,
@@ -1844,7 +1852,13 @@ class FlashAttentionBackend(AttentionBackend):
aux_tensors=None,
rel_bias=None,
rel_bias_event=None,
) -> torch.Tensor:
# Returns (output, lse) with lse in [total_q, num_heads].
return_lse: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
lse_out = None
# Bound in __init__ so a subclass can substitute a different FA4 build.
flash_attn_with_kvcache = self.flash_attn_with_kvcache
flash_attn_varlen_func = self.flash_attn_varlen_func
if score_mod is not None and self.fa_impl_ver != 4:
raise RuntimeError("score_mod is only supported by the FA4 backend.")
if k is not None:
@@ -2045,7 +2059,7 @@ class FlashAttentionBackend(AttentionBackend):
causal=False if use_cascade_attn else causal,
window_size=window_size,
softcap=layer.logit_cap,
return_softmax_lse=use_cascade_attn,
return_softmax_lse=use_cascade_attn or return_lse,
num_splits=(
self.decode_num_splits
if not is_swa_layer
@@ -2090,6 +2104,8 @@ class FlashAttentionBackend(AttentionBackend):
o_expand,
softmax_lse_expand.T.contiguous(),
)
elif return_lse:
o, lse_out, *_ = result
else:
o = result
else:
@@ -2168,7 +2184,12 @@ class FlashAttentionBackend(AttentionBackend):
else:
o = result
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
o = o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
if return_lse:
assert lse_out is not None
# The varlen kernel emits LSE head-major [num_heads, total_q].
return o, lse_out.transpose(0, 1).contiguous()
return o
def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
"""Initialize CUDA graph state for the attention backend.