Uploaded using `kernel-builder`.

Browse files

Files changed (5) hide show

build/torch212-cxx11-cu132-x86_64-linux/__init__.py +16 -2
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu132-x86_64-linux/{_sageattention3_blackwell_cuda_464a940.abi3.so → _sageattention3_blackwell_cuda_f7be481.abi3.so} +1 -1
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +7 -7
build/torch212-cxx11-cu132-x86_64-linux/sageattention3_blackwell/__init__.py +0 -26

build/torch212-cxx11-cu132-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/torch212-cxx11-cu132-x86_64-linux/_ops.py CHANGED Viewed

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

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

build/torch212-cxx11-cu132-x86_64-linux/{_sageattention3_blackwell_cuda_464a940.abi3.so → _sageattention3_blackwell_cuda_f7be481.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:65b634001e530983df71b273e081a3b95fb1967257cc1b2efc242d672740a4ff
 size 4576288

build/torch212-cxx11-cu132-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,9 +1,10 @@
 {
   "name": "sageattention3-blackwell",
-  "id": "_sageattention3_blackwell_cuda_464a940",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
@@ -13,20 +14,19 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "kKV+PdmaMj3xE7oZtkaAiPO2w08YCeGkWqledC8C1XQ=",
-      "_ops.py": "o8pimywpFhJ/C6mmzD9jYF2h1V7p0pRtf417UEO7x9U=",
-      "_sageattention3_blackwell_cuda_464a940.abi3.so": "2bMjNT1QjjBNTTWg6O/mGWZFvnbfUgpD56pGMiq7vsQ=",
-      "sageattention3_blackwell/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
   "provenance": {
     "kernel-builder": {
       "version": "0.17.0-dev0",
-      "sha": "b39ca23f1b36383df00b27b3ffe1276cd5dbea85",
       "dirty": false
     },
     "kernel": {
-      "sha": "464a94081ba465f65e25b14019320c3910c305c6",
       "dirty": false
     }
   }

 {
   "name": "sageattention3-blackwell",
+  "id": "_sageattention3_blackwell_cuda_f7be481",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
+  "kernel-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "fFNOV1L7fwXJ/GAKEPWGjndpGeObnj+C6saq6qdBkC8=",
+      "_ops.py": "H7CUYyZh3jmz/iSVRJddN7iHH0e/iHt+fYV4vmCBqDw=",
+      "_sageattention3_blackwell_cuda_f7be481.abi3.so": "ZbY0AB5TCYPfcbJz4IGjuV+xlnJXzBsu/CQtZydApP8="
     }
   },
   "provenance": {
     "kernel-builder": {
       "version": "0.17.0-dev0",
+      "sha": "3a5c09d941a19f582fe551d038dc2b3bc6b0a03d",
       "dirty": false
     },
     "kernel": {
+      "sha": "f7be4812381108c475322303ae4638273bc566fd",
       "dirty": false
     }
   }

build/torch212-cxx11-cu132-x86_64-linux/sageattention3_blackwell/__init__.py DELETED Viewed

@@ -1,26 +0,0 @@
-import ctypes
-import importlib.util
-import sys
-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")))