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

Upgrade fa3 hash (#28394)

Browse Source 
Co-authored-by: Shijin <dovis.zhang02@gmail.com>
This commit is contained in:
Baizhou Zhang
2026-06-17 13:32:11 -07:00
committed by GitHub
co-authored by Shijin
parent 753aa89a83
commit 4b817f5d7f
6 changed files with 82 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/flash_attention.py
+2
View File
@@ -35,6 +35,7 @@ def flash_attn_with_kvcache(
scheduler_metadata=None,
num_splits=0, # Can be tuned for speed
pack_gqa=None, # Can be tuned for speed
only_qv=False, # ver=3 only: skip K matmul when qk rope dim is 0
sm_margin=0, # Can be tuned if some SMs are used for communication
return_softmax_lse=False,
sinks=None,
@@ -162,6 +163,7 @@ def flash_attn_with_kvcache(
scheduler_metadata=scheduler_metadata,
num_splits=num_splits,
pack_gqa=pack_gqa,
only_qv=only_qv,
sm_margin=sm_margin,
return_softmax_lse=return_softmax_lse,
sinks=sinks,
python/sglang/jit_kernel/flash_attention_v3.py
+12 -4
View File
@@ -128,6 +128,7 @@ def flash_attn_with_kvcache(
scheduler_metadata=None,
num_splits=0, # Can be tuned for speed
pack_gqa=None, # Can be tuned for speed
only_qv=False, # Skip K matmul when qk rope dim is 0 (requires qv)
sm_margin=0, # Can be tuned if some SMs are used for communication
return_softmax_lse=False,
sinks=None,
@@ -138,7 +139,11 @@ def flash_attn_with_kvcache(
"flash_attn at sgl-kernel is only supported on sm90 and above"
)
assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
# When only_qv=True the caller may pass k_cache=None (synthetic K is
# allocated inside the sgl-kernel wrapper). Skip the stride check in that
# case so the rope=0 path doesn't trip the assertion.
if k_cache is not None:
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"
return _call_fa3_kernel(
@@ -171,9 +176,10 @@ def flash_attn_with_kvcache(
scheduler_metadata,
num_splits,
pack_gqa,
sm_margin,
return_softmax_lse,
sinks,
sm_margin=sm_margin,
only_qv=only_qv,
return_softmax_lse=return_softmax_lse,
sinks=sinks,
out=out,
)
@@ -201,6 +207,7 @@ def flash_attn_varlen_func(
softcap=0.0,
num_splits=1,
pack_gqa=None,
only_qv=False,
sm_margin=0,
return_softmax_lse=False,
sinks=None,
@@ -265,6 +272,7 @@ def flash_attn_varlen_func(
softcap=softcap,
num_splits=num_splits,
pack_gqa=pack_gqa,
only_qv=only_qv,
sm_margin=sm_margin,
return_softmax_lse=return_softmax_lse,
sinks=sinks,
sgl-kernel/CMakeLists.txt
+12 -5
View File
@@ -81,8 +81,8 @@ FetchContent_Populate(repo-flashinfer)
# flash-attention
FetchContent_Declare(
repo-flash-attention
URL https://${GITHUB_ARTIFACTORY}/sgl-project/sgl-attn/archive/bcf72ccc6816b36a5fae2c5a3c027604629785e0.tar.gz
URL_HASH SHA256=2110d8ca1ed9b330b9f99c1d4088be12a37e44152c39fdf90dfb6526f532ba96
URL https://${GITHUB_ARTIFACTORY}/sgl-project/sgl-attn/archive/65c54cc5a6d29fee56036484c749bc5b8e00fd66.tar.gz
URL_HASH SHA256=53307eefba65d9ac26092433a8bc4d48182ce0bb27132e69ab0b8d6ff8ee47db
)
FetchContent_Populate(repo-flash-attention)
@@ -105,6 +105,14 @@ endif()
option(ENABLE_BELOW_SM90 "Enable gencode below SM90" ${DEFAULT_ENABLE_BELOW_SM90})
set(DEFAULT_SGL_KERNEL_ENABLE_FA3 OFF)
if (CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$")
message(STATUS "For aarch64, disable FA3 by default")
endif()
if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.4" AND NOT CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$")
set(DEFAULT_SGL_KERNEL_ENABLE_FA3 ON)
endif()
include_directories(
${PROJECT_SOURCE_DIR}/include
${PROJECT_SOURCE_DIR}/csrc
@@ -165,7 +173,7 @@ list(APPEND SGL_KERNEL_CUDA_FLAGS
option(SGL_KERNEL_ENABLE_BF16 "Enable BF16" ON)
option(SGL_KERNEL_ENABLE_FP8 "Enable FP8" ON)
option(SGL_KERNEL_ENABLE_FP4 "Enable FP4" OFF)
option(SGL_KERNEL_ENABLE_FA3 "Enable FA3" OFF)
option(SGL_KERNEL_ENABLE_FA3 "Enable FA3" ${DEFAULT_SGL_KERNEL_ENABLE_FA3})
option(SGL_KERNEL_ENABLE_SM90A "Enable SM90A" OFF)
option(SGL_KERNEL_ENABLE_SM100A "Enable SM100A" OFF)
@@ -222,8 +230,7 @@ if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_SM100A)
endif()
endif()
if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.4")
set(SGL_KERNEL_ENABLE_FA3 ON)
if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.4" AND SGL_KERNEL_ENABLE_FA3)
list(APPEND SGL_KERNEL_CUDA_FLAGS
"-gencode=arch=compute_90a,code=sm_90a"
)
sgl-kernel/csrc/flash_extension.cc
+3 -1
View File
@@ -57,7 +57,9 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
" int num_splits,"
" bool? pack_gqa,"
" int sm_margin,"
" Tensor? sinks"
" Tensor? sinks,"
" Tensor? sparse_mask_fine," // [total_q, max_k_blocks, num_int32_per_block]
" bool only_qv"
") -> (Tensor(a!), Tensor, Tensor, Tensor)"); // first return aliases out
m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));
sgl-kernel/include/sgl_flash_kernel_ops.h
+3 -1
View File
@@ -82,7 +82,9 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_fwd(
int64_t num_splits,
std::optional<bool> pack_gqa_,
int64_t sm_margin,
std::optional<const at::Tensor>& sinks_); // (h)
std::optional<const at::Tensor>& sinks_, // (h)
std::optional<at::Tensor> sparse_mask_fine_, // [total_q, max_k_blocks, num_int32_per_block]
bool only_qv);
/*
* From flash-attention: get_scheduler_metadata
sgl-kernel/python/sgl_kernel/flash_attn.py
+50 -5
View File
@@ -62,6 +62,7 @@ def flash_attn_with_kvcache(
scheduler_metadata=None,
num_splits=0, # Can be tuned for speed
pack_gqa=None, # Can be tuned for speed
only_qv=False, # Only use QV (skip K matmul); requires qv. Used when qk rope dim is 0.
sm_margin=0, # Can be tuned if some SMs are used for communication
return_softmax_lse=False,
sinks=None,
@@ -159,15 +160,54 @@ def flash_attn_with_kvcache(
normalization factor).
"""
assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
if v_cache is None:
raise ValueError("v_cache must be provided")
assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension"
if k_cache is None:
if not only_qv:
raise ValueError("k_cache can only be None when only_qv=True")
if q is not None:
k_head_size = q.shape[-1]
k_dtype = q.dtype
k_device = q.device
elif k is not None:
k_head_size = k.shape[-1]
k_dtype = k.dtype
k_device = k.device
else:
# Fallback: only_qv kernel ignores K values, so a tiny placeholder works.
k_head_size = 64
k_dtype = v_cache.dtype
k_device = v_cache.device
k_shape = (*v_cache.shape[:-1], k_head_size)
# The kernel path for only_qv ignores K values, but backend API still requires k tensor.
k_cache = torch.empty(k_shape, dtype=k_dtype, device=k_device)
assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
if q is None:
if not only_qv:
raise ValueError("q can only be None when only_qv=True")
if qv is None:
raise ValueError(
"q must be provided unless qv is provided with only_qv=True"
)
q_shape = (*qv.shape[:-1], k_cache.shape[-1])
# The kernel path for only_qv ignores q values, but backend API still requires q tensor.
q = torch.empty(q_shape, dtype=qv.dtype, device=qv.device)
if softmax_scale is None:
softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (
-0.5
)
if only_qv:
if qv is None:
raise ValueError("only_qv=True requires qv to be provided")
softmax_scale = (qv.shape[-1]) ** (-0.5)
else:
softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (
-0.5
)
if cache_seqlens is not None and isinstance(cache_seqlens, int):
cache_seqlens = torch.full(
(k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device
(q.shape[0],), cache_seqlens, dtype=torch.int32, device=v_cache.device
)
cache_seqlens = maybe_contiguous(cache_seqlens)
@@ -223,6 +263,8 @@ def flash_attn_with_kvcache(
pack_gqa,
sm_margin,
sinks,
None, # sparse_mask_fine
only_qv,
)
# return (out, softmax_lse) if return_softmax_lse else out
return (out, softmax_lse, *rest) if return_softmax_lse else out
@@ -251,6 +293,7 @@ def flash_attn_varlen_func(
softcap=0.0,
num_splits=1,
pack_gqa=None,
only_qv=False,
sm_margin=0,
return_softmax_lse=False,
sinks=None,
@@ -311,6 +354,8 @@ def flash_attn_varlen_func(
pack_gqa=pack_gqa,
sm_margin=sm_margin,
sinks=sinks,
sparse_mask_fine=None,
only_qv=only_qv,
)
return (out, softmax_lse, *rest) if return_softmax_lse else out