minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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_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_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( def flash_attn_with_kvcache(
@@ -168,6 +166,10 @@ def flash_attn_with_kvcache(
sinks=sinks, sinks=sinks,
) )
elif ver == 4: elif ver == 4:
from .flash_attention_v4 import (
flash_attn_with_kvcache as fa4_flash_attn_with_kvcache,
)
return fa4_flash_attn_with_kvcache( return fa4_flash_attn_with_kvcache(
q, q,
k_cache, k_cache,
@@ -260,6 +262,10 @@ def flash_attn_varlen_func(
sinks=sinks, sinks=sinks,
) )
elif ver == 4: elif ver == 4:
from .flash_attention_v4 import (
flash_attn_varlen_func as fa4_flash_attn_varlen_func,
)
return fa4_flash_attn_varlen_func( return fa4_flash_attn_varlen_func(
q, q,
k, k,
python/sglang/srt/model_loader/weight_utils.py
+2 -2
View File
@@ -551,9 +551,9 @@ def download_safetensors_index_file_from_hf(
# If file not found on remote or locally, we should not fail since # If file not found on remote or locally, we should not fail since
# only some models will have index_file. # only some models will have index_file.
except huggingface_hub.utils.EntryNotFoundError: except huggingface_hub.utils.EntryNotFoundError:
logger.info("No %s found in remote.", index_file) logger.debug("No %s found in remote.", index_file)
except huggingface_hub.utils.LocalEntryNotFoundError: except huggingface_hub.utils.LocalEntryNotFoundError:
logger.info("No %s found in local cache.", index_file) logger.debug("No %s found in local cache.", index_file)
# For models like Mistral-7B-v0.3, there are both sharded # For models like Mistral-7B-v0.3, there are both sharded
python/sglang/srt/utils/numa_utils.py
+2 -2
View File
@@ -56,7 +56,7 @@ def _mp_set_executable(executable: str, debug_str: str):
old_executable = os.fsdecode(multiprocessing.spawn.get_executable()) old_executable = os.fsdecode(multiprocessing.spawn.get_executable())
multiprocessing.spawn.set_executable(executable) multiprocessing.spawn.set_executable(executable)
logger.info(f"mp.set_executable {old_executable} -> {executable} ({debug_str})") logger.debug(f"mp.set_executable {old_executable} -> {executable} ({debug_str})")
try: try:
yield yield
finally: finally:
@@ -64,7 +64,7 @@ def _mp_set_executable(executable: str, debug_str: str):
os.fsdecode(multiprocessing.spawn.get_executable()) == executable os.fsdecode(multiprocessing.spawn.get_executable()) == executable
), f"{multiprocessing.spawn.get_executable()=}" ), f"{multiprocessing.spawn.get_executable()=}"
multiprocessing.spawn.set_executable(old_executable) multiprocessing.spawn.set_executable(old_executable)
logger.info(f"mp.set_executable revert to {old_executable}") logger.debug(f"mp.set_executable revert to {old_executable}")
def get_numa_node_if_available(server_args: ServerArgs, gpu_id: int) -> Optional[int]: def get_numa_node_if_available(server_args: ServerArgs, gpu_id: int) -> Optional[int]: