diff --git a/python/sglang/jit_kernel/flash_attention.py b/python/sglang/jit_kernel/flash_attention.py index da6c59dae..c6a67757d 100644 --- a/python/sglang/jit_kernel/flash_attention.py +++ b/python/sglang/jit_kernel/flash_attention.py @@ -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, diff --git a/python/sglang/jit_kernel/flash_attention_v3.py b/python/sglang/jit_kernel/flash_attention_v3.py index 78423daaa..fe7f42234 100644 --- a/python/sglang/jit_kernel/flash_attention_v3.py +++ b/python/sglang/jit_kernel/flash_attention_v3.py @@ -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, diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index 1f376d535..99015ff18 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -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" ) diff --git a/sgl-kernel/csrc/flash_extension.cc b/sgl-kernel/csrc/flash_extension.cc index 1fd387f6a..d466d936d 100644 --- a/sgl-kernel/csrc/flash_extension.cc +++ b/sgl-kernel/csrc/flash_extension.cc @@ -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)); diff --git a/sgl-kernel/include/sgl_flash_kernel_ops.h b/sgl-kernel/include/sgl_flash_kernel_ops.h index 10ee93075..271095482 100644 --- a/sgl-kernel/include/sgl_flash_kernel_ops.h +++ b/sgl-kernel/include/sgl_flash_kernel_ops.h @@ -82,7 +82,9 @@ std::tuple mha_fwd( int64_t num_splits, std::optional pack_gqa_, int64_t sm_margin, - std::optional& sinks_); // (h) + std::optional& sinks_, // (h) + std::optional sparse_mask_fine_, // [total_q, max_k_blocks, num_int32_per_block] + bool only_qv); /* * From flash-attention: get_scheduler_metadata diff --git a/sgl-kernel/python/sgl_kernel/flash_attn.py b/sgl-kernel/python/sgl_kernel/flash_attn.py index e498f1967..635171b6c 100644 --- a/sgl-kernel/python/sgl_kernel/flash_attn.py +++ b/sgl-kernel/python/sgl_kernel/flash_attn.py @@ -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