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

Lazy import flash_attention_v4 to avoid loading flash_attn.cute at startup (#22306)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-04-08 20:40:25 -07:00
committed by GitHub
parent f127d67823
commit ddc8ef1038
3 changed files with 12 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/flash_attention.py
+8 -2
View File
@@ -4,8 +4,6 @@ import torch
from .flash_attention_v3 import flash_attn_varlen_func as fa3_flash_attn_varlen_func
from .flash_attention_v3 import flash_attn_with_kvcache as fa3_flash_attn_with_kvcache
from .flash_attention_v4 import flash_attn_varlen_func as fa4_flash_attn_varlen_func
from .flash_attention_v4 import flash_attn_with_kvcache as fa4_flash_attn_with_kvcache
def flash_attn_with_kvcache(
@@ -168,6 +166,10 @@ def flash_attn_with_kvcache(
sinks=sinks,
)
elif ver == 4:
from .flash_attention_v4 import (
flash_attn_with_kvcache as fa4_flash_attn_with_kvcache,
)
return fa4_flash_attn_with_kvcache(
q,
k_cache,
@@ -260,6 +262,10 @@ def flash_attn_varlen_func(
sinks=sinks,
)
elif ver == 4:
from .flash_attention_v4 import (
flash_attn_varlen_func as fa4_flash_attn_varlen_func,
)
return fa4_flash_attn_varlen_func(
q,
k,