Build uploaded using `kernels`.

.gitattributes

@@ -37,3 +37,8 @@ build/torch210-cxx11-cu128-x86_64-linux/_sgl_flash_attn3_cuda_b8d1001.abi3.so fi
 build/torch210-cxx11-cu130-x86_64-linux/_sgl_flash_attn3_cuda_b8d1001.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch29-cxx11-cu128-x86_64-linux/_sgl_flash_attn3_cuda_b8d1001.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch29-cxx11-cu130-x86_64-linux/_sgl_flash_attn3_cuda_b8d1001.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu128-x86_64-linux/_sgl_flash_attn3_cuda_1988a8e.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu130-x86_64-linux/_sgl_flash_attn3_cuda_1988a8e.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-cxx11-cu128-x86_64-linux/_sgl_flash_attn3_cuda_1988a8e.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-cxx11-cu129-x86_64-linux/_sgl_flash_attn3_cuda_1988a8e.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-cxx11-cu130-x86_64-linux/_sgl_flash_attn3_cuda_1988a8e.abi3.so filter=lfs diff=lfs merge=lfs -text

build/torch210-cxx11-cu128-x86_64-linux/__init__.py

@@ -86,7 +86,7 @@ def flash_attn_with_kvcache(
     rotary_seqlens = maybe_contiguous(rotary_seqlens)
     attention_chunk = 0 if attention_chunk is None else int(attention_chunk)
 
-    out, softmax_lse, *rest = ops.fwd(
+    out, softmax_lse, *_ = ops.fwd(
         q,
         k_cache,
         v_cache,
@@ -123,7 +123,8 @@ def flash_attn_with_kvcache(
         sm_margin,
         sinks,
     )
-    return (out, softmax_lse, *rest) if return_softmax_lse else out
+
+    return (out, softmax_lse) if return_softmax_lse else out
 
 
 def flash_attn_varlen_func(
@@ -166,7 +167,7 @@ def flash_attn_varlen_func(
         )
     attention_chunk = 0 if attention_chunk is None else int(attention_chunk)
 
-    out, softmax_lse, *rest = ops.fwd(
+    out, softmax_lse, *_ = ops.fwd(
         q,
         k,
         v,
@@ -204,4 +205,4 @@ def flash_attn_varlen_func(
         sinks,
     )
 
-    return (out, softmax_lse, *rest) if return_softmax_lse else out
+    return (out, softmax_lse) if return_softmax_lse else out

build/torch210-cxx11-cu128-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _sgl_flash_attn3_cuda_b8d1001
-ops = torch.ops._sgl_flash_attn3_cuda_b8d1001
+from . import _sgl_flash_attn3_cuda_1988a8e
+ops = torch.ops._sgl_flash_attn3_cuda_1988a8e
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_sgl_flash_attn3_cuda_b8d1001::{op_name}"
+    return f"_sgl_flash_attn3_cuda_1988a8e::{op_name}"

build/torch210-cxx11-cu128-x86_64-linux/{_sgl_flash_attn3_cuda_b8d1001.abi3.so → _sgl_flash_attn3_cuda_1988a8e.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:12d0b5a344ad2d758ed9170172db410b3cb272bea0a80d13fb694bf19f843ec9
+oid sha256:e6c2c92bea67e95646f7ede11a3a8a683c95467b988344155708fca7329358fc
 size 916761672

build/torch210-cxx11-cu128-x86_64-linux/metadata.json

@@ -1,5 +1,12 @@
 {
   "version": 1,
   "license": "BSD-3-Clause",
-  "python-depends": []
-}
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "9.0a"
+    ]
+  }
+}

build/torch210-cxx11-cu130-x86_64-linux/__init__.py

@@ -86,7 +86,7 @@ def flash_attn_with_kvcache(
     rotary_seqlens = maybe_contiguous(rotary_seqlens)
     attention_chunk = 0 if attention_chunk is None else int(attention_chunk)
 
-    out, softmax_lse, *rest = ops.fwd(
+    out, softmax_lse, *_ = ops.fwd(
         q,
         k_cache,
         v_cache,
@@ -123,7 +123,8 @@ def flash_attn_with_kvcache(
         sm_margin,
         sinks,
     )
-    return (out, softmax_lse, *rest) if return_softmax_lse else out
+
+    return (out, softmax_lse) if return_softmax_lse else out
 
 
 def flash_attn_varlen_func(
@@ -166,7 +167,7 @@ def flash_attn_varlen_func(
         )
     attention_chunk = 0 if attention_chunk is None else int(attention_chunk)
 
-    out, softmax_lse, *rest = ops.fwd(
+    out, softmax_lse, *_ = ops.fwd(
         q,
         k,
         v,
@@ -204,4 +205,4 @@ def flash_attn_varlen_func(
         sinks,
     )
 
-    return (out, softmax_lse, *rest) if return_softmax_lse else out
+    return (out, softmax_lse) if return_softmax_lse else out

build/torch210-cxx11-cu130-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _sgl_flash_attn3_cuda_b8d1001
-ops = torch.ops._sgl_flash_attn3_cuda_b8d1001
+from . import _sgl_flash_attn3_cuda_1988a8e
+ops = torch.ops._sgl_flash_attn3_cuda_1988a8e
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_sgl_flash_attn3_cuda_b8d1001::{op_name}"
+    return f"_sgl_flash_attn3_cuda_1988a8e::{op_name}"

build/torch210-cxx11-cu130-x86_64-linux/{_sgl_flash_attn3_cuda_b8d1001.abi3.so → _sgl_flash_attn3_cuda_1988a8e.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:84e5bf1bc8feb0ca30e06b3ece38d0659ad55844de0de0a389e246209c658cad
+oid sha256:763d2698ae560844d581a83f0ce4e8da584849f1beeb2e8ebf008fdc3db10c7a
 size 917921992

build/torch210-cxx11-cu130-x86_64-linux/metadata.json

@@ -1,5 +1,12 @@
 {
   "version": 1,
   "license": "BSD-3-Clause",
-  "python-depends": []
-}
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "9.0a"
+    ]
+  }
+}

build/torch29-cxx11-cu128-x86_64-linux/__init__.py

@@ -86,7 +86,7 @@ def flash_attn_with_kvcache(
     rotary_seqlens = maybe_contiguous(rotary_seqlens)
     attention_chunk = 0 if attention_chunk is None else int(attention_chunk)
 
-    out, softmax_lse, *rest = ops.fwd(
+    out, softmax_lse, *_ = ops.fwd(
         q,
         k_cache,
         v_cache,
@@ -123,7 +123,8 @@ def flash_attn_with_kvcache(
         sm_margin,
         sinks,
     )
-    return (out, softmax_lse, *rest) if return_softmax_lse else out
+
+    return (out, softmax_lse) if return_softmax_lse else out
 
 
 def flash_attn_varlen_func(
@@ -166,7 +167,7 @@ def flash_attn_varlen_func(
         )
     attention_chunk = 0 if attention_chunk is None else int(attention_chunk)
 
-    out, softmax_lse, *rest = ops.fwd(
+    out, softmax_lse, *_ = ops.fwd(
         q,
         k,
         v,
@@ -204,4 +205,4 @@ def flash_attn_varlen_func(
         sinks,
     )
 
-    return (out, softmax_lse, *rest) if return_softmax_lse else out
+    return (out, softmax_lse) if return_softmax_lse else out

build/torch29-cxx11-cu128-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _sgl_flash_attn3_cuda_b8d1001
-ops = torch.ops._sgl_flash_attn3_cuda_b8d1001
+from . import _sgl_flash_attn3_cuda_1988a8e
+ops = torch.ops._sgl_flash_attn3_cuda_1988a8e
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_sgl_flash_attn3_cuda_b8d1001::{op_name}"
+    return f"_sgl_flash_attn3_cuda_1988a8e::{op_name}"

build/torch29-cxx11-cu128-x86_64-linux/{_sgl_flash_attn3_cuda_b8d1001.abi3.so → _sgl_flash_attn3_cuda_1988a8e.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:105ceff789a1d4ec5a32b4000a0b09cadd89d2d8aeace62f0fe860b495acba56
+oid sha256:918b5b761bc5150da21a2af42ea8c8892fc51ef1200a4febaa1a05d706400ddb
 size 916755168

build/torch29-cxx11-cu128-x86_64-linux/metadata.json

@@ -1,5 +1,12 @@
 {
   "version": 1,
   "license": "BSD-3-Clause",
-  "python-depends": []
-}
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "9.0a"
+    ]
+  }
+}

build/torch29-cxx11-cu129-x86_64-linux/__init__.py

@@ -0,0 +1,208 @@
+from functools import lru_cache
+from typing import Optional, Union
+
+import torch
+
+from ._ops import ops
+
+
+@lru_cache(maxsize=1)
+def is_fa3_supported(device=None) -> bool:
+    return (torch.version.cuda >= "12.3") and (
+        torch.cuda.get_device_capability(device)[0] == 9
+        or torch.cuda.get_device_capability(device)[0] == 8
+    )
+
+
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+
+
+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),
+    attention_chunk: Optional[int] = None,
+    softcap=0.0,
+    rotary_interleaved=True,
+    scheduler_metadata=None,
+    num_splits=0,
+    pack_gqa=None,
+    sm_margin=0,
+    return_softmax_lse=False,
+    sinks=None,
+):
+    """
+    FA3 flash_attn_with_kvcache: forward-only attention with paged KV cache,
+    optional rotary embedding, sliding window, softcapping, GQA/MQA, and
+    inplace KV cache update.
+
+    See sgl-kernel documentation for full argument descriptions.
+    """
+    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"
+    if softmax_scale is None:
+        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
+        )
+        cache_seqlens = maybe_contiguous(cache_seqlens)
+
+    q, k_cache, k, v = [maybe_contiguous(x) for x in (q, k_cache, k, v)]
+    v_cache = (
+        v_cache.contiguous()
+        if v_cache.stride(-1) != 1 and v_cache.stride(-3) != 1
+        else v_cache
+    )
+    cu_seqlens_q, cu_seqlens_k_new = [
+        maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k_new)
+    ]
+    page_table, cache_batch_idx, cache_leftpad = [
+        maybe_contiguous(x) for x in (page_table, cache_batch_idx, cache_leftpad)
+    ]
+    rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)]
+    rotary_seqlens = maybe_contiguous(rotary_seqlens)
+    attention_chunk = 0 if attention_chunk is None else int(attention_chunk)
+
+    out, softmax_lse, *_ = ops.fwd(
+        q,
+        k_cache,
+        v_cache,
+        k,
+        v,
+        qv,
+        None,  # out
+        cu_seqlens_q,
+        None,  # cu_seqlens_k
+        cu_seqlens_k_new,
+        None,  # seqused_q
+        cache_seqlens,
+        max_seqlen_q,
+        None,  # max_seqlen_k
+        page_table,
+        cache_batch_idx,
+        cache_leftpad,
+        rotary_cos,
+        rotary_sin,
+        rotary_seqlens,
+        q_descale,
+        k_descale,
+        v_descale,
+        softmax_scale,
+        causal,
+        window_size[0],
+        window_size[1],
+        attention_chunk,
+        softcap,
+        rotary_interleaved,
+        scheduler_metadata,
+        num_splits,
+        pack_gqa,
+        sm_margin,
+        sinks,
+    )
+
+    return (out, softmax_lse) if return_softmax_lse else out
+
+
+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(
+            "sgl_flash_attn3 is only supported on sm80 and above with CUDA >= 12.3"
+        )
+
+    if max_seqlen_q is None or max_seqlen_k is None:
+        raise ValueError("max_seqlen_q and max_seqlen_k are required for FA3")
+
+    if softmax_scale is None:
+        softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (
+            -0.5
+        )
+    attention_chunk = 0 if attention_chunk is None else int(attention_chunk)
+
+    out, softmax_lse, *_ = ops.fwd(
+        q,
+        k,
+        v,
+        None,  # k_new
+        None,  # v_new
+        qv,
+        None,  # out
+        cu_seqlens_q,
+        cu_seqlens_k,
+        None,  # cu_seqlens_k_new
+        seqused_q,
+        seqused_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        None,  # page_table
+        None,  # kv_batch_idx
+        None,  # leftpad_k
+        None,  # rotary cos
+        None,  # rotary sin
+        None,  # seqlens_rotary
+        q_descale,
+        k_descale,
+        v_descale,
+        softmax_scale,
+        causal,
+        window_size[0],
+        window_size[1],
+        attention_chunk,
+        softcap,
+        False,  # is_rotary_interleaved
+        None,  # scheduler_metadata
+        num_splits,
+        pack_gqa,
+        sm_margin,
+        sinks,
+    )
+
+    return (out, softmax_lse) if return_softmax_lse else out

build/torch29-cxx11-cu129-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _sgl_flash_attn3_cuda_1988a8e
+ops = torch.ops._sgl_flash_attn3_cuda_1988a8e
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_sgl_flash_attn3_cuda_1988a8e::{op_name}"

build/torch29-cxx11-cu129-x86_64-linux/_sgl_flash_attn3_cuda_1988a8e.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f40e01d849dfe7c460d1ff7c3a53396ce293024b241002cee8842ac8c6cb01be
+size 921853920

build/torch29-cxx11-cu129-x86_64-linux/metadata.json

@@ -0,0 +1,12 @@
+{
+  "version": 1,
+  "license": "BSD-3-Clause",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "9.0a"
+    ]
+  }
+}

build/torch29-cxx11-cu129-x86_64-linux/sgl_flash_attn3/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+from pathlib import Path
+from types import ModuleType
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch29-cxx11-cu130-x86_64-linux/__init__.py

@@ -86,7 +86,7 @@ def flash_attn_with_kvcache(
     rotary_seqlens = maybe_contiguous(rotary_seqlens)
     attention_chunk = 0 if attention_chunk is None else int(attention_chunk)
 
-    out, softmax_lse, *rest = ops.fwd(
+    out, softmax_lse, *_ = ops.fwd(
         q,
         k_cache,
         v_cache,
@@ -123,7 +123,8 @@ def flash_attn_with_kvcache(
         sm_margin,
         sinks,
     )
-    return (out, softmax_lse, *rest) if return_softmax_lse else out
+
+    return (out, softmax_lse) if return_softmax_lse else out
 
 
 def flash_attn_varlen_func(
@@ -166,7 +167,7 @@ def flash_attn_varlen_func(
         )
     attention_chunk = 0 if attention_chunk is None else int(attention_chunk)
 
-    out, softmax_lse, *rest = ops.fwd(
+    out, softmax_lse, *_ = ops.fwd(
         q,
         k,
         v,
@@ -204,4 +205,4 @@ def flash_attn_varlen_func(
         sinks,
     )
 
-    return (out, softmax_lse, *rest) if return_softmax_lse else out
+    return (out, softmax_lse) if return_softmax_lse else out

build/torch29-cxx11-cu130-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _sgl_flash_attn3_cuda_b8d1001
-ops = torch.ops._sgl_flash_attn3_cuda_b8d1001
+from . import _sgl_flash_attn3_cuda_1988a8e
+ops = torch.ops._sgl_flash_attn3_cuda_1988a8e
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_sgl_flash_attn3_cuda_b8d1001::{op_name}"
+    return f"_sgl_flash_attn3_cuda_1988a8e::{op_name}"

build/torch29-cxx11-cu130-x86_64-linux/{_sgl_flash_attn3_cuda_b8d1001.abi3.so → _sgl_flash_attn3_cuda_1988a8e.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5055457995b929e955171991186babcd494f435feba5d37dc21c25ec805f98cd
+oid sha256:c23184d7250226dad7f0d023593f9c45b9db57c9f1daa6e08e216a32f78df8fa
 size 917907296

build/torch29-cxx11-cu130-x86_64-linux/metadata.json

@@ -1,5 +1,12 @@
 {
   "version": 1,
   "license": "BSD-3-Clause",
-  "python-depends": []
-}
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "9.0a"
+    ]
+  }
+}