sgl-project--sglang/python/sglang/jit_kernel/flash_attention_v3.py

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
from sglang.srt.utils import get_device_capability, is_musa

logger = logging.getLogger(__name__)

SGL_FA3_KERNEL_REPO = "kernels-community/sgl-flash-attn3"
SGL_FA3_KERNEL_REVISION = "v1"
DEFAULT_FA3_KERNEL_LOCKFILE = "kernels.lock"


def _call_fa3_kernel(kernel, *args, out=None, **kwargs):
    if out is None:
        return kernel(*args, **kwargs)
    try:
        return kernel(*args, **kwargs, out=out)
    except TypeError as exc:
        if "unexpected keyword argument 'out'" not in str(exc):
            raise
        return kernel(*args, **kwargs)


@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.
    major, minor = get_device_capability()
    if is_musa():
        return major >= 3
    if torch.version.cuda is not None and torch.version.cuda >= "12.3":
        return major == 9 or major == 8
    return False


@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
    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,
    out=None,
):
    if not _is_fa3_supported():
        raise NotImplementedError(
            "flash_attn at sgl-kernel is only supported on sm90 and above"
        )

    # 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(
        _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=sm_margin,
        only_qv=only_qv,
        return_softmax_lse=return_softmax_lse,
        sinks=sinks,
        out=out,
    )


@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,
    only_qv=False,
    sm_margin=0,
    return_softmax_lse=False,
    sinks=None,
    out=None,
):

    if not _is_fa3_supported():
        # Fall back to flash_attn package (FA2) on platforms without sgl-kernel FA3
        # (e.g. ROCm, or CUDA < sm90)
        if cu_seqlens_q is not None:
            from flash_attn import flash_attn_varlen_func as fa2_flash_attn_varlen_func

            return fa2_flash_attn_varlen_func(
                q,
                k,
                v,
                cu_seqlens_q,
                cu_seqlens_k,
                max_seqlen_q,
                max_seqlen_k,
                softmax_scale=softmax_scale,
                causal=causal,
                window_size=window_size,
                softcap=softcap,
                return_attn_probs=return_softmax_lse,
            )
        else:
            # 4D inputs (batch, seqlen, nheads, headdim) without cu_seqlens
            from flash_attn import flash_attn_func as fa2_flash_attn_func

            return fa2_flash_attn_func(
                q,
                k,
                v,
                softmax_scale=softmax_scale,
                causal=causal,
                window_size=window_size,
                softcap=softcap,
                return_attn_probs=return_softmax_lse,
            )

    return _call_fa3_kernel(
        _load_fa3_kernels()["flash_attn_varlen_func"],
        q=q,
        k=k,
        v=v,
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_k=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,
        only_qv=only_qv,
        sm_margin=sm_margin,
        return_softmax_lse=return_softmax_lse,
        sinks=sinks,
        out=out,
    )