Uploaded using `kernel-builder`.

Browse files

Files changed (5) hide show

benchmarks/benchmark.py +6 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +16 -2
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu128-x86_64-linux/{_sageattention3_blackwell_cuda_77467d6_dirty.abi3.so → _sageattention3_blackwell_cuda_2dc0021.abi3.so} +2 -2
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +6 -6

benchmarks/benchmark.py CHANGED Viewed

@@ -63,9 +63,15 @@ def main() -> None:
             trust_remote_code=True,
         )
     cases = [(6144, 128), (2688, 64)]
     if args.mode == "full":
         cases = [(6144, 128), (24576, 128), (2688, 64)]
     print("| S | D | SDPA us | Sage2 static us | Sage3 core+quant us | Sage3 fused eager us | Sage3 fused graph us | graph vs SDPA | fused/legacy cosine |")
     print("|---:|---:|---:|---:|---:|---:|---:|---:|---:|")
     for s, d in cases:

             trust_remote_code=True,
         )
+    supported_head_dims = (
+        tuple(ops.module.capabilities()["head_dims"])
+        if args.backend == "installed" else (64, 128)
+    )
     cases = [(6144, 128), (2688, 64)]
     if args.mode == "full":
         cases = [(6144, 128), (24576, 128), (2688, 64)]
+    cases = [(s, d) for s, d in cases if d in supported_head_dims]
     print("| S | D | SDPA us | Sage2 static us | Sage3 core+quant us | Sage3 fused eager us | Sage3 fused graph us | graph vs SDPA | fused/legacy cosine |")
     print("|---:|---:|---:|---:|---:|---:|---:|---:|---:|")
     for s, d in cases:

build/torch211-cxx11-cu128-x86_64-linux/__init__.py CHANGED Viewed

@@ -10,7 +10,16 @@ import torch
 from ._ops import add_op_namespace_prefix, ops
-SUPPORTED_HEAD_DIMS = (64, 128)
 SUPPORTED_LAYOUTS = ("NHD",)
 TOKEN_ALIGNMENT = 128
 ACCURACY_PROFILE = "speed-first"
@@ -29,6 +38,7 @@ def capabilities() -> dict[str, object]:
         "cuda_graph_safe": True,
         "fused_prep": True,
         "delta_dtypes": ("float32", "bfloat16"),
     }
@@ -62,7 +72,11 @@ class Sage3FusedWorkspace:
 def _check_nhd(x: torch.Tensor, name: str) -> None:
     if x.dim() != 4 or x.shape[-1] not in SUPPORTED_HEAD_DIMS:
-        raise RuntimeError(f"{name} must have contiguous NHD shape [B,L,H,64|128]")
     if not x.is_cuda or not x.is_contiguous():
         raise RuntimeError(f"{name} must be contiguous CUDA")
     if x.dtype not in (torch.float16, torch.bfloat16):

 from ._ops import add_op_namespace_prefix, ops
+def _cuda_version_tuple() -> tuple[int, int]:
+    version = torch.version.cuda
+    if not version:
+        return (0, 0)
+    major, minor, *_ = version.split(".")
+    return (int(major), int(minor))
+CUDA_VERSION = _cuda_version_tuple()
+SUPPORTED_HEAD_DIMS = (64, 128) if CUDA_VERSION >= (13, 0) else (64,)
 SUPPORTED_LAYOUTS = ("NHD",)
 TOKEN_ALIGNMENT = 128
 ACCURACY_PROFILE = "speed-first"
         "cuda_graph_safe": True,
         "fused_prep": True,
         "delta_dtypes": ("float32", "bfloat16"),
+        "d128_min_cuda": "13.0",
     }
 def _check_nhd(x: torch.Tensor, name: str) -> None:
     if x.dim() != 4 or x.shape[-1] not in SUPPORTED_HEAD_DIMS:
+        supported = "|".join(str(dim) for dim in SUPPORTED_HEAD_DIMS)
+        raise RuntimeError(
+            f"{name} must have contiguous NHD shape [B,L,H,{supported}] "
+            f"for this CUDA {torch.version.cuda} artifact"
+        )
     if not x.is_cuda or not x.is_contiguous():
         raise RuntimeError(f"{name} must be contiguous CUDA")
     if x.dtype not in (torch.float16, torch.bfloat16):

build/torch211-cxx11-cu128-x86_64-linux/_ops.py CHANGED Viewed

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

 import torch
+from . import _sageattention3_blackwell_cuda_2dc0021
+ops = torch.ops._sageattention3_blackwell_cuda_2dc0021
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_sageattention3_blackwell_cuda_2dc0021::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/{_sageattention3_blackwell_cuda_77467d6_dirty.abi3.so → _sageattention3_blackwell_cuda_2dc0021.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:025175c02c4682c09ca50df5cdaa814b1207fbf9974c1250e74063bbfd96a21f
-size 4542576

 version https://git-lfs.github.com/spec/v1
+oid sha256:b63380dfb4f0692a87dfc2edb38951b0606fd28c17842314ac24ec9c7b9f7c14
+size 2304832

build/torch211-cxx11-cu128-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "sageattention3-blackwell",
-  "id": "_sageattention3_blackwell_cuda_77467d6_dirty",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,9 +13,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "kKV+PdmaMj3xE7oZtkaAiPO2w08YCeGkWqledC8C1XQ=",
-      "_ops.py": "dgBUJ1px3wMMLuiDDeoNj5FWj26ENaYLiQ7O+yG7ZHo=",
-      "_sageattention3_blackwell_cuda_77467d6_dirty.abi3.so": "AlF1wCxGgsCcpQ31zaqBSxIH+/mXTBJQ50Bju/2Woh8=",
       "sageattention3_blackwell/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -26,8 +26,8 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "77467d6151612cde2f6aef1364feaa06bd65f639",
-      "dirty": true
     }
   }
 }

 {
   "name": "sageattention3-blackwell",
+  "id": "_sageattention3_blackwell_cuda_2dc0021",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "fFNOV1L7fwXJ/GAKEPWGjndpGeObnj+C6saq6qdBkC8=",
+      "_ops.py": "/0FdyIGUwzozcipWgUccGqq+FM6yEsqSB3H1hewc4fI=",
+      "_sageattention3_blackwell_cuda_2dc0021.abi3.so": "tjOA37TwaSqH38Lts4lRsGBv0owXhCMUrCTsnHuffBQ=",
       "sageattention3_blackwell/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
       "dirty": false
     },
     "kernel": {
+      "sha": "2dc0021cd0c6cb1e4c5258662a17051385215bd6",
+      "dirty": false
     }
   }
 }