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

Move fa4 from sgl-kernel to jit kernel (#17353)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-01-24 15:25:03 +08:00
committed by GitHub
parent 7a4bb0d516
commit 3992a023e6
45 changed files with 24594 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
python/sglang/srt/layers/attention/flashattention_backend.py
View File

@@ -22,7 +22,18 @@ if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
from sgl_kernel import merge_state_v2
from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache
from sgl_kernel.flash_attn import flash_attn_varlen_func as flash_attn_varlen_func_fa3
from sgl_kernel.flash_attn import flash_attn_with_kvcache as flash_attn_with_kvcache_fa3
flash_attn_varlen_func = flash_attn_varlen_func_fa3
flash_attn_with_kvcache = flash_attn_with_kvcache_fa3
from sglang.jit_kernel.flash_attention_v4 import (
flash_attn_varlen_func as flash_attn_varlen_func_fa4,
)
from sglang.jit_kernel.flash_attention_v4 import (
flash_attn_with_kvcache as flash_attn_with_kvcache_fa4,
)
@dataclass
@@ -803,10 +814,21 @@ class FlashAttentionBackend(AttentionBackend):
and not is_swa_layer
)
# For fa3 interface version compatibility, we put new fields into conditional keyword args
flash_attn_varlen_func_base = flash_attn_varlen_func_fa3
flash_attn_with_kvcache_base = flash_attn_with_kvcache_fa3
flash_attn_varlen_func = (
flash_attn_varlen_func_fa4
if self.fa_impl_ver == 4
else flash_attn_varlen_func_base
)
flash_attn_with_kvcache = (
flash_attn_with_kvcache_fa4
if self.fa_impl_ver == 4
else flash_attn_with_kvcache_base
)
kwargs = {}
if self.fa_impl_ver != 3:
kwargs["ver"] = self.fa_impl_ver
if sinks is not None:
kwargs["sinks"] = sinks
@@ -1110,10 +1132,7 @@ class FlashAttentionBackend(AttentionBackend):
if layer.is_cross_attention or layer.attn_type == AttentionType.ENCODER_ONLY:
causal = False
# For fa3 interface version compatibility, we put new fields into conditional keyword args
kwargs = {}
if self.fa_impl_ver != 3:
kwargs["ver"] = self.fa_impl_ver
if sinks is not None:
kwargs["sinks"] = sinks