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

Kernels community fa3 (#20796)

Browse Source
This commit is contained in:
Rain Jiang
2026-04-07 12:48:44 -07:00
committed by GitHub
parent 0c204fbd57
commit 1a8eb890f6
20 changed files with 1956 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/pyproject.toml
+4
View File
@@ -77,6 +77,7 @@ dependencies = [
"watchfiles",
"xgrammar==0.1.32",
"smg-grpc-servicer>=0.5.0",
"kernels",
]
[[tool.uv.index]]
@@ -207,3 +208,6 @@ version_file = "sglang/_version.py"
git_describe_command = ["python3", "python/tools/get_version_tag.py", "--tag-only"]
# Allow editable installs even when .git metadata is not available.
fallback_version = "0.0.0.dev0"
[tool.kernels.dependencies]
"kernels-community/sgl-flash-attn3" = 1
python/sglang/jit_kernel/flash_attention.py
+286
View File
@@ -0,0 +1,286 @@
from typing import Optional, Union
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(
q,
k_cache,
v_cache,
k=None,
v=None,
qv=None,
rotary_cos=None,
rotary_sin=None,
cache_seqlens: Optional[Union[int, torch.Tensor]] = None,
cache_batch_idx: Optional[torch.Tensor] = None,
cache_leftpad: Optional[torch.Tensor] = None,
page_table: Optional[torch.Tensor] = None,
cu_seqlens_q: Optional[torch.Tensor] = None,
cu_seqlens_k_new: Optional[torch.Tensor] = None,
max_seqlen_q: Optional[int] = None,
rotary_seqlens: Optional[torch.Tensor] = None,
q_descale: Optional[torch.Tensor] = None,
k_descale: Optional[torch.Tensor] = None,
v_descale: Optional[torch.Tensor] = None,
softmax_scale=None,
causal=False,
window_size=(-1, -1), # -1 means infinite context window
attention_chunk: Optional[int] = None,
softcap=0.0, # 0.0 means deactivated
rotary_interleaved=True,
scheduler_metadata=None,
num_splits=0, # Can be tuned for speed
pack_gqa=None, # Can be tuned for speed
sm_margin=0, # Can be tuned if some SMs are used for communication
return_softmax_lse=False,
sinks=None,
score_mod=None,
aux_tensors=None,
ver=3,
):
"""
If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
k and v. This is useful for incremental decoding: you can pass in the cached keys/values from
the previous step, and update them with the new keys/values from the current step, and do
attention with the updated cache, all in 1 kernel.
If you pass in k / v, you must make sure that the cache is large enough to hold the new values.
For example, the KV cache could be pre-allocated with the max sequence length, and you can use
cache_seqlens to keep track of the current sequence lengths of each sequence in the batch.
Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be
rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos
and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at
indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).
See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.
Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
1 1 1 1 0
1 1 1 1 1
If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
0 0
0 0
0 0
1 0
1 1
If the row of the mask is all zero, the output will be zero.
If window_size != (-1, -1), implements sliding window local attention. Query at position i
will only attend to keys between
[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
Note: Does not support backward pass.
Arguments:
q: (batch_size, seqlen, nheads, headdim)
k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table,
or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache)
page_block_size must be a multiple of 256.
v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table,
or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache)
k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate
k with k_cache, starting at the indices specified by cache_seqlens.
v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k.
qv [optional]: (batch_size, seqlen, nheads, headdim_v)
rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding
to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.
rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.
cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the
KV cache.
cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.
If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].
If the indices are not distinct, and k and v are provided, the values updated in the cache
might come from any of the duplicate indices.
cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0.
page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.
softmax_scale: float. The scaling of QK^T before applying softmax.
Default to 1 / sqrt(headdim).
causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
window_size: (left, right). If not (-1, -1), implements sliding window local attention.
attention_chunk: Optional[int]. If not None, splits the query into chunks of this size to save memory.
softcap: float. Anything > 0 activates softcapping attention.
rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.
If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,
rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1
(i.e. GPT-NeoX style).
num_splits: int. If > 1, split the key/value into this many chunks along the sequence.
If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic
to automatically determine the number of splits.
Don't change this unless you know what you are doing.
return_softmax_lse: bool. Whether to return the logsumexp of the attention scores.
score_mod [optional]: A callable that takes the attention scores and applies a modification.
aux_tensors [optional]: Some score_mods will want to read from global aux_tensors. This is how we thread them through to the inner kernel.
Return:
out: (batch_size, seqlen, nheads, headdim).
softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The
logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
normalization factor).
"""
if ver == 3:
return fa3_flash_attn_with_kvcache(
q,
k_cache,
v_cache,
k=k,
v=v,
qv=qv,
rotary_cos=rotary_cos,
rotary_sin=rotary_sin,
cache_seqlens=cache_seqlens,
cache_batch_idx=cache_batch_idx,
cache_leftpad=cache_leftpad,
page_table=page_table,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k_new=cu_seqlens_k_new,
max_seqlen_q=max_seqlen_q,
rotary_seqlens=rotary_seqlens,
q_descale=q_descale,
k_descale=k_descale,
v_descale=v_descale,
softmax_scale=softmax_scale,
causal=causal,
window_size=window_size,
attention_chunk=attention_chunk,
softcap=softcap,
rotary_interleaved=rotary_interleaved,
scheduler_metadata=scheduler_metadata,
num_splits=num_splits,
pack_gqa=pack_gqa,
sm_margin=sm_margin,
return_softmax_lse=return_softmax_lse,
sinks=sinks,
)
elif ver == 4:
return fa4_flash_attn_with_kvcache(
q,
k_cache,
v_cache,
k=k,
v=v,
qv=qv,
rotary_cos=rotary_cos,
rotary_sin=rotary_sin,
cache_seqlens=cache_seqlens,
cache_batch_idx=cache_batch_idx,
cache_leftpad=cache_leftpad,
page_table=page_table,
cu_seqlens_q=cu_seqlens_q,
max_seqlen_q=max_seqlen_q,
rotary_seqlens=rotary_seqlens,
q_descale=q_descale,
k_descale=k_descale,
v_descale=v_descale,
softmax_scale=softmax_scale,
causal=causal,
window_size=window_size,
softcap=softcap,
num_splits=num_splits,
pack_gqa=pack_gqa,
sinks=sinks,
score_mod=score_mod,
aux_tensors=aux_tensors,
return_softmax_lse=return_softmax_lse,
)
else:
raise RuntimeError(f"Unknown flash attention version {ver}")
def flash_attn_varlen_func(
q,
k,
v,
cu_seqlens_q,
cu_seqlens_k,
max_seqlen_q=None,
max_seqlen_k=None,
seqused_q=None,
seqused_k=None,
page_table=None,
softmax_scale=None,
causal=False,
qv=None,
q_descale=None,
k_descale=None,
v_descale=None,
window_size=(-1, -1),
attention_chunk=0,
softcap=0.0,
num_splits=1,
pack_gqa=None,
sm_margin=0,
return_softmax_lse=False,
sinks=None,
score_mod=None,
aux_tensors=None,
ver=3,
):
if ver == 3:
return fa3_flash_attn_varlen_func(
q,
k,
v,
cu_seqlens_q,
cu_seqlens_k,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
seqused_q=seqused_q,
seqused_k=seqused_k,
page_table=page_table,
softmax_scale=softmax_scale,
causal=causal,
qv=qv,
q_descale=q_descale,
k_descale=k_descale,
v_descale=v_descale,
window_size=window_size,
attention_chunk=attention_chunk,
softcap=softcap,
num_splits=num_splits,
pack_gqa=pack_gqa,
sm_margin=sm_margin,
return_softmax_lse=return_softmax_lse,
sinks=sinks,
)
elif ver == 4:
return fa4_flash_attn_varlen_func(
q,
k,
v,
cu_seqlens_q,
cu_seqlens_k,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
seqused_q=seqused_q,
seqused_k=seqused_k,
page_table=page_table,
softmax_scale=softmax_scale,
causal=causal,
softcap=softcap,
window_size=window_size,
sinks=sinks,
num_splits=num_splits,
pack_gqa=pack_gqa,
score_mod=score_mod,
aux_tensors=aux_tensors,
return_softmax_lse=return_softmax_lse,
)
else:
raise RuntimeError(f"Unknown flash attention version {ver}")
python/sglang/jit_kernel/flash_attention_v3.py
+222
View File
@@ -0,0 +1,222 @@
import logging
import os
from typing import Optional, Union
import torch
from sglang.jit_kernel.utils import cache_once
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.environ import envs
logger = logging.getLogger(__name__)
SGL_FA3_KERNEL_REPO = "kernels-community/sgl-flash-attn3"
SGL_FA3_KERNEL_REVISION = "v1"
DEFAULT_FA3_KERNEL_LOCKFILE = "kernels.lock"
@cache_once
def _load_fa3_kernels():
# By default, we use the implementation from sgl-kernel,
# which is expected to be more stable and compatible
if envs.SGLANG_USE_SGL_FA3_KERNEL.get():
logger.debug(
f"SGLANG_USE_SGL_FA3_KERNEL=True, use sgl-kernel implementation for FlashAttention v3 "
)
return _load_fa3_kernel_from_sgl()
# Otherwise, we try to load the kernels from the kernels community cache directory or kernels community repo
lockfile_path = os.path.join(
envs.SGLANG_CACHE_DIR.get(), DEFAULT_FA3_KERNEL_LOCKFILE
)
try:
from kernels import get_kernel, load_kernel
# When the lock file provided, load from the kernel cache directory,
# otherwise, load from the repo, which require download from huggingface hub
# but always works as long as the repo is accessible.
if os.path.exists(lockfile_path):
ops = load_kernel(SGL_FA3_KERNEL_REPO, lockfile_path)
else:
ops = get_kernel(SGL_FA3_KERNEL_REPO, revision=SGL_FA3_KERNEL_REVISION)
return {
"flash_attn_with_kvcache": ops.flash_attn_with_kvcache,
"flash_attn_varlen_func": ops.flash_attn_varlen_func,
}
except Exception as e:
# When the kernels from the repo or the cache directory cannot be loaded
# we catch the exception and log a warning, and then fallback to the implementation
# from sgl-kernel, which is expected to be less efficient but more compatible.
logger.warning(
f"Rollback to implementation from sgl-kernel since loading FlashAttention v3 "
f"kernels from {SGL_FA3_KERNEL_REPO} with lockfile {lockfile_path} failed: {e}"
)
return _load_fa3_kernel_from_sgl()
def _load_fa3_kernel_from_sgl():
from sgl_kernel.flash_attn import (
flash_attn_varlen_func,
flash_attn_with_kvcache,
)
return {
"flash_attn_with_kvcache": flash_attn_with_kvcache,
"flash_attn_varlen_func": flash_attn_varlen_func,
}
@cache_once
def _is_fa3_supported(device=None) -> bool:
# There some fa3 FYI
# FA3 can fail without a enough shared memory for a some shapes, such as higher
# hidden_dim or some special cases.
# Right now, fa3 is supported for sm80/sm87 and sm86/sm89. The main different
# Between sm80/sm87 and sm86/sm89 is the shared memory size. you can follow the link below for more information
# https://docs.nvidia.com/cuda/cuda-c-programming-guide/#shared-memory-8-x
# And for sgl-kernel right now, we can build fa3 on sm80/sm86/sm89/sm90a.
# That means if you use A100/A*0/L20/L40/L40s/4090 you can use fa3.
return (torch.version.cuda >= "12.3") and (
torch.cuda.get_device_capability(device)[0] == 9
or torch.cuda.get_device_capability(device)[0] == 8
)
@debug_kernel_api
def flash_attn_with_kvcache(
q,
k_cache,
v_cache,
k=None,
v=None,
qv=None,
rotary_cos=None,
rotary_sin=None,
cache_seqlens: Optional[Union[int, torch.Tensor]] = None,
cache_batch_idx: Optional[torch.Tensor] = None,
cache_leftpad: Optional[torch.Tensor] = None,
page_table: Optional[torch.Tensor] = None,
cu_seqlens_q: Optional[torch.Tensor] = None,
cu_seqlens_k_new: Optional[torch.Tensor] = None,
max_seqlen_q: Optional[int] = None,
rotary_seqlens: Optional[torch.Tensor] = None,
q_descale: Optional[torch.Tensor] = None,
k_descale: Optional[torch.Tensor] = None,
v_descale: Optional[torch.Tensor] = None,
softmax_scale=None,
causal=False,
window_size=(-1, -1), # -1 means infinite context window
attention_chunk: Optional[int] = None,
softcap=0.0, # 0.0 means deactivated
rotary_interleaved=True,
scheduler_metadata=None,
num_splits=0, # Can be tuned for speed
pack_gqa=None, # Can be tuned for speed
sm_margin=0, # Can be tuned if some SMs are used for communication
return_softmax_lse=False,
sinks=None,
):
if not _is_fa3_supported():
raise NotImplementedError(
"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"
assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension"
return _load_fa3_kernels()["flash_attn_with_kvcache"](
q,
k_cache,
v_cache,
k,
v,
qv,
rotary_cos,
rotary_sin,
cache_seqlens,
cache_batch_idx,
cache_leftpad,
page_table,
cu_seqlens_q,
cu_seqlens_k_new,
max_seqlen_q,
rotary_seqlens,
q_descale,
k_descale,
v_descale,
softmax_scale,
causal,
window_size,
attention_chunk,
softcap,
rotary_interleaved,
scheduler_metadata,
num_splits,
pack_gqa,
sm_margin,
return_softmax_lse,
sinks,
)
@debug_kernel_api
def flash_attn_varlen_func(
q,
k,
v,
cu_seqlens_q,
cu_seqlens_k,
max_seqlen_q=None,
max_seqlen_k=None,
seqused_q=None,
seqused_k=None,
page_table=None,
softmax_scale=None,
causal=False,
qv=None,
q_descale=None,
k_descale=None,
v_descale=None,
window_size=(-1, -1),
attention_chunk=0,
softcap=0.0,
num_splits=1,
pack_gqa=None,
sm_margin=0,
return_softmax_lse=False,
sinks=None,
):
if not _is_fa3_supported():
raise NotImplementedError(
"flash_attn at sgl-kernel is only supported on sm90 and above"
)
return _load_fa3_kernels()["flash_attn_varlen_func"](
q,
k,
v,
cu_seqlens_q,
cu_seqlens_k,
max_seqlen_q,
max_seqlen_k,
seqused_q,
seqused_k,
page_table,
softmax_scale,
causal,
qv,
q_descale,
k_descale,
v_descale,
window_size,
attention_chunk,
softcap,
num_splits,
pack_gqa,
sm_margin,
return_softmax_lse,
sinks,
)
python/sglang/jit_kernel/flash_attention_v4.py
-1
View File
@@ -42,7 +42,6 @@ def flash_attn_varlen_func(
score_mod: Optional[Callable] = None,
aux_tensors: Optional[list] = None,
return_softmax_lse: bool = False,
**_: object,
):
if _flash_attn_varlen_func is None: # pragma: no cover
raise ImportError(
python/sglang/jit_kernel/tests/test_flash_attention_3.py
+1373
View File
File diff suppressed because it is too large Load Diff
python/sglang/jit_kernel/tests/test_flash_attention_4.py
+3 -1
View File
@@ -11,7 +11,7 @@ import torch
import torch.nn.functional as F
from einops import rearrange, repeat
from sglang.jit_kernel.flash_attention_v4 import flash_attn_varlen_func
from sglang.jit_kernel.flash_attention import flash_attn_varlen_func
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=120, suite="stage-b-kernel-unit-1-gpu-large")
@@ -826,6 +826,7 @@ def test_flash_attn_varlen_output(
sinks=learnable_sink, # FA4 uses learnable_sink, not sinks
pack_gqa=pack_gqa,
return_softmax_lse=True,
ver=4,
)
out = output_pad_fn(out_unpad)
if query_unused_mask is not None:
@@ -1384,6 +1385,7 @@ def test_flash_attn_kvcache(
softcap=0.0,
pack_gqa=None,
return_softmax_lse=True,
ver=4,
)
if varlen_q:
out = output_pad_fn(out)
python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
+4 -18
View File
@@ -5,27 +5,13 @@ from typing import Any, List, Optional, Tuple
import torch
from sglang.jit_kernel.flash_attention import flash_attn_varlen_func
from sglang.multimodal_gen.runtime.layers.utils import register_custom_op
from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
from sglang.multimodal_gen.runtime.platforms import (
AttentionBackendEnum,
)
try:
from sgl_kernel.flash_attn import flash_attn_varlen_func
from sglang.jit_kernel.flash_attention_v4 import (
flash_attn_varlen_func as flash_attn_varlen_func_fa4,
)
def flash_attn_func(*args, ver: int = 3, **kwargs):
if ver == 4:
return flash_attn_varlen_func_fa4(*args, **kwargs)
return flash_attn_varlen_func(*args, **kwargs)
except ImportError as e:
raise e
def maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
return x.contiguous() if x is not None and x.stride(-1) != 1 else x
@@ -207,7 +193,7 @@ def flash_attn_varlen_func_op(
"flash_attn_varlen_func_op is out-only op; return_softmax_lse must be False. "
"Use flash_attn_varlen_func_op_lse for (out, lse)."
)
return flash_attn_func(
return flash_attn_varlen_func(
q,
k,
v,
@@ -271,7 +257,7 @@ def flash_attn_varlen_func_op_lse(
"flash_attn_varlen_func_op_lse is out+lse op; return_softmax_lse must be True. "
"Use flash_attn_varlen_func_op for out-only."
)
return flash_attn_func(
return flash_attn_varlen_func(
q,
k,
v,
@@ -409,7 +395,7 @@ class FlashAttentionImpl(AttentionImpl):
# - fa_ver == 3: call python function (can return Tensor or (Tensor, Tensor) depending on flag)
# - fa_ver == 4: call custom ops with FIXED return schema
if fa_ver == 3:
flash_attn_op = flash_attn_func
flash_attn_op = flash_attn_varlen_func
output = flash_attn_op(
q=query,
k=key,
python/sglang/srt/compilation/backend.py
+2 -3
View File
@@ -21,6 +21,7 @@ from sglang.srt.compilation.compiler_interface import EagerAdapter, InductorAdap
from sglang.srt.compilation.cuda_piecewise_backend import CUDAPiecewiseBackend
from sglang.srt.compilation.npu_piecewise_backend import NPUPiecewiseBackend
from sglang.srt.compilation.pass_manager import PostGradPassManager
from sglang.srt.environ import envs
from sglang.srt.utils.common import is_npu
logger = logging.getLogger(__name__)
@@ -393,9 +394,7 @@ class SGLangBackend:
self.inductor_config["post_grad_custom_post_pass"] = self.post_grad_pass_manager
def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable:
base_cache_dir = os.path.expanduser(
os.getenv("SGLANG_CACHE_DIR", "~/.cache/sglang/")
)
base_cache_dir = envs.SGLANG_CACHE_DIR.get()
cache_hash = self.compiler_manager.compute_hash()
cache_dir = os.path.join(
python/sglang/srt/environ.py
+6
View File
@@ -406,6 +406,9 @@ class Envs:
# sgl-kernel
SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)
# Flash Attention
SGLANG_USE_SGL_FA3_KERNEL = EnvBool(True)
# vLLM dependencies (TODO: they have been deprecated, we can remove them safely)
USE_VLLM_CUTLASS_W8A8_FP8_KERNEL = EnvBool(False)
@@ -534,6 +537,9 @@ class Envs:
# Elastic EP Backup Port
SGLANG_BACKUP_PORT_BASE = EnvInt(10000)
# Sglang Cache Dir
SGLANG_CACHE_DIR = EnvStr(os.path.expanduser("~/.cache/sglang"))
envs = Envs()
EnvField._allow_set_name = False
python/sglang/srt/layers/attention/dual_chunk_flashattention_backend.py
+4 -1
View File
@@ -9,13 +9,16 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
import torch
import torch.nn.functional as F
from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache
from sgl_kernel.sparse_flash_attn import (
convert_vertical_slash_indexes,
convert_vertical_slash_indexes_mergehead,
sparse_attn_func,
)
from sglang.jit_kernel.flash_attention import (
flash_attn_varlen_func,
flash_attn_with_kvcache,
)
from sglang.srt.distributed.parallel_state import get_tensor_model_parallel_rank
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.flashattention_backend import FlashAttentionMetadata
python/sglang/srt/layers/attention/flashattention_backend.py
+19 -6
View File
@@ -27,6 +27,11 @@ if TYPE_CHECKING:
from sgl_kernel import merge_state_v2
from sglang.jit_kernel.flash_attention import (
flash_attn_varlen_func,
flash_attn_with_kvcache,
)
@dataclass
class FlashAttentionMetadata:
@@ -616,9 +621,6 @@ class FlashAttentionBackend(AttentionBackend):
and not is_swa_layer
)
flash_attn_varlen_func = self.flash_attn_varlen_func
flash_attn_with_kvcache = self.flash_attn_with_kvcache
kwargs = {}
if sinks is not None:
kwargs["sinks"] = sinks
@@ -696,6 +698,7 @@ class FlashAttentionBackend(AttentionBackend):
v_descale=v_descale,
return_softmax_lse=use_cascade_attn,
num_splits=self.num_splits,
ver=self.fa_impl_ver,
**kwargs,
)
@@ -723,6 +726,7 @@ class FlashAttentionBackend(AttentionBackend):
v_descale=v_descale,
return_softmax_lse=use_cascade_attn,
num_splits=self.num_splits,
ver=self.fa_impl_ver,
**kwargs,
)
@@ -750,6 +754,7 @@ class FlashAttentionBackend(AttentionBackend):
v_descale=v_descale,
return_softmax_lse=True,
num_splits=self.num_splits,
ver=self.fa_impl_ver,
**kwargs,
)
o, _ = merge_state_v2_wrapper(
@@ -789,6 +794,7 @@ class FlashAttentionBackend(AttentionBackend):
softmax_scale=layer.scaling,
causal=False,
return_softmax_lse=True,
ver=self.fa_impl_ver,
**kwargs,
)
else:
@@ -814,6 +820,7 @@ class FlashAttentionBackend(AttentionBackend):
softmax_scale=layer.scaling,
causal=True,
return_softmax_lse=forward_batch.mha_return_lse,
ver=self.fa_impl_ver,
**kwargs,
)
if forward_batch.mha_return_lse:
@@ -822,7 +829,7 @@ class FlashAttentionBackend(AttentionBackend):
return output, lse
return output
else:
assert self.fa_impl_ver in [3], "Only FA3 support here"
assert self.fa_impl_ver == 3, "Only FA3 support here"
# Do absorbed multi-latent attention
kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(
layer.layer_id
@@ -865,6 +872,7 @@ class FlashAttentionBackend(AttentionBackend):
v_descale=v_descale,
return_softmax_lse=use_cascade_attn,
num_splits=self.num_splits,
ver=self.fa_impl_ver,
)
if use_cascade_attn:
o, softmax_lse, *rest = result
@@ -887,6 +895,7 @@ class FlashAttentionBackend(AttentionBackend):
v_descale=v_descale,
return_softmax_lse=True,
num_splits=self.num_splits,
ver=self.fa_impl_ver,
)
)
o, _ = merge_state_v2_wrapper(
@@ -964,8 +973,6 @@ class FlashAttentionBackend(AttentionBackend):
if sinks is not None:
kwargs["sinks"] = sinks
flash_attn_with_kvcache = self.flash_attn_with_kvcache
k_descale, v_descale = None, None
# only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention
# has corresponding quantization method so that layer.k_scale is not None,
@@ -1009,6 +1016,7 @@ class FlashAttentionBackend(AttentionBackend):
k_descale=k_descale,
v_descale=v_descale,
num_splits=self.num_splits,
ver=self.fa_impl_ver,
**kwargs,
)
elif use_local_attn:
@@ -1029,6 +1037,7 @@ class FlashAttentionBackend(AttentionBackend):
k_descale=k_descale,
v_descale=v_descale,
num_splits=self.num_splits,
ver=self.fa_impl_ver,
**kwargs,
)
else:
@@ -1066,6 +1075,7 @@ class FlashAttentionBackend(AttentionBackend):
v_descale=v_descale,
return_softmax_lse=use_cascade_attn,
num_splits=self.num_splits,
ver=self.fa_impl_ver,
**kwargs,
)
if use_cascade_attn:
@@ -1088,6 +1098,7 @@ class FlashAttentionBackend(AttentionBackend):
v_descale=v_descale,
return_softmax_lse=True,
num_splits=self.num_splits,
ver=self.fa_impl_ver,
**kwargs,
)
)
@@ -1144,6 +1155,7 @@ class FlashAttentionBackend(AttentionBackend):
v_descale=v_descale,
return_softmax_lse=use_cascade_attn, # softmax_lse is needed for merge states
num_splits=self.num_splits,
ver=self.fa_impl_ver,
)
if use_cascade_attn:
o, softmax_lse, *rest = result
@@ -1165,6 +1177,7 @@ class FlashAttentionBackend(AttentionBackend):
v_descale=v_descale,
return_softmax_lse=True,
num_splits=self.num_splits,
ver=self.fa_impl_ver,
)
o, _ = merge_state_v2(
o,
python/sglang/srt/layers/attention/nsa_backend.py
+4 -1
View File
@@ -61,7 +61,10 @@ if _is_hip:
"aiter is AMD specific kernel library. Please make sure aiter is installed on your AMD device."
)
else:
from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache
from sglang.jit_kernel.flash_attention import (
flash_attn_varlen_func,
flash_attn_with_kvcache,
)
# Reuse this workspace buffer across all NSA backend instances
python/sglang/srt/layers/attention/vision.py
+6 -18
View File
@@ -38,21 +38,9 @@ _is_xpu = is_xpu()
if _is_cuda:
from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
try:
from sgl_kernel.flash_attn import flash_attn_varlen_func
def flash_attn_func(*args, ver: int = 3, **kwargs):
if ver == 4:
from sglang.jit_kernel.flash_attention_v4 import (
flash_attn_varlen_func as flash_attn_varlen_func_fa4,
)
return flash_attn_varlen_func_fa4(*args, **kwargs)
return flash_attn_varlen_func(*args, **kwargs)
except ImportError as e:
raise e
from sglang.jit_kernel.flash_attention import (
flash_attn_varlen_func,
)
if _is_npu:
import torch_npu
@@ -420,7 +408,7 @@ class VisionFlash3Attention(nn.Module):
"""
if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
max_seqlen = cu_seqlens[1]
output = flash_attn_func(
output = flash_attn_varlen_func(
q,
k,
v,
@@ -436,7 +424,7 @@ class VisionFlash3Attention(nn.Module):
seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
max_seqlen = seq_lens.max().item()
output = flash_attn_func(
output = flash_attn_varlen_func(
q,
k,
v,
@@ -489,7 +477,7 @@ class VisionFlash4Attention(nn.Module):
seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
max_seqlen = seq_lens.max().item()
output = flash_attn_func(
output = flash_attn_varlen_func(
q,
k,
v,
python/sglang/srt/layers/attention/xpu_backend.py
+5 -1
View File
@@ -20,7 +20,11 @@ if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
from sgl_kernel import merge_state_v2
from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache
from sglang.jit_kernel.flash_attention import (
flash_attn_varlen_func,
flash_attn_with_kvcache,
)
class XPUAttentionBackend(AttentionBackend):
python/sglang/srt/utils/runai_utils.py
+3 -7
View File
@@ -5,6 +5,8 @@ import logging
import os
from pathlib import Path
from sglang.srt.environ import envs
logger = logging.getLogger(__name__)
SUPPORTED_SCHEMES = ["s3://", "gs://", "az://"]
@@ -26,12 +28,6 @@ SUPPORTED_SCHEMES = ["s3://", "gs://", "az://"]
# This avoids file locks, race conditions, and duplicate downloads
def get_cache_dir() -> str:
# Expand user path (~) to ensure absolute paths for locking
path = os.getenv("SGLANG_CACHE_DIR", "~/.cache/sglang/")
return os.path.expanduser(path)
def list_safetensors(path: str = "") -> list[str]:
"""
List full file names from object path and filter by allow pattern.
@@ -122,7 +118,7 @@ class ObjectStorageModel:
Returns the local directory path.
"""
model_hash = hashlib.sha256(str(model_path).encode()).hexdigest()[:16]
base_dir = get_cache_dir()
base_dir = envs.SGLANG_CACHE_DIR.get()
# Ensure base cache dir exists
os.makedirs(os.path.join(base_dir, "model_streamer"), exist_ok=True)