EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

fix: update flash attn (#5308)

Browse Source
This commit is contained in:
Yineng Zhang
2025-04-11 16:23:09 -07:00
committed by GitHub
parent 136b8e6afb
commit 4f288113ce
1 changed files with 4 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
sgl-kernel/python/sgl_kernel/flash_attn.py
View File

@@ -10,15 +10,9 @@ except:
def is_fa3_supported(device=None) -> bool: def is_fa3_supported(device=None) -> bool:
# FA3 can fail without a enough shared memory for a some shapes, currently # now sgl-kernel only build fa3 for sm90a && cuda >= 12.3
# only 8.0 and 8.7 have enough shared memory for all shapes return (torch.cuda.get_device_capability(device)[0] == 9) and (
# https://docs.nvidia.com/cuda/cuda-c-programming-guide/#shared-memory-8-x torch.version.cuda >= "12.3"
# now sgl-kernel only build fa3 for sm90a && cuda >= 12.4
return (
(torch.cuda.get_device_capability(device)[0] >= 9)
and (torch.version.cuda >= "12.4")
# or torch.cuda.get_device_capability(device) == (8, 0)
# or torch.cuda.get_device_capability(device) == (8, 7)
) )
@@ -144,7 +138,7 @@ def flash_attn_with_kvcache(
""" """
if not is_fa3_supported(): if not is_fa3_supported():
raise NotImplementedError( raise NotImplementedError(
"flash_attn at sgl-kernel is only supported on sm90 and above" "flash_attn at sgl-kernel is only supported on sm90 and cu123 above"
) )
assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension"