Uploaded using `kernel-builder`.

Files changed (4) hide show

build/torch212-cxx11-cu130-x86_64-linux/__init__.py CHANGED Viewed

@@ -11,8 +11,9 @@ import torch
 from ._ops import add_op_namespace_prefix, ops
-SUPPORTED_HEAD_DIMS = (64, 96, 128, 256)
-SPLIT_HEAD_DIMS = (96, 128, 256)
 @dataclass(frozen=True)
@@ -43,7 +44,7 @@ def recommended_num_splits(
     if any(int(v) <= 0 for v in values):
         raise ValueError("all shape values and num_sms must be positive")
     if int(head_dim) not in SUPPORTED_HEAD_DIMS:
-        raise ValueError(f"head_dim must be one of {SUPPORTED_HEAD_DIMS}")
     block_n = 256 if head_dim <= 64 else (128 if head_dim <= 128 else 64)
     n_blocks = _ceildiv(seqlen_k, block_n)
     m_blocks = _ceildiv(seqlen_q, 64)
@@ -82,8 +83,8 @@ def allocate_workspace(
     if q.ndim != 4 or k.ndim != 4:
         raise ValueError("q and k must have shape (B, S, H, D)")
-    if q.shape[-1] not in SPLIT_HEAD_DIMS:
-        return None
     if num_sms is None:
         num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count
     splits = recommended_num_splits(
@@ -96,8 +97,9 @@ def allocate_workspace(
         device=q.device,
         dtype=torch.float32,
     )
     out = torch.empty(
-        (splits, q.shape[0], q.shape[2], q.shape[1], q.shape[3]),
         device=q.device,
         dtype=torch.float32,
     )
@@ -273,6 +275,7 @@ def forward(
 __all__ = [
     "FA2Workspace",
     "SPLIT_HEAD_DIMS",
     "SUPPORTED_HEAD_DIMS",
     "allocate_outputs",

 from ._ops import add_op_namespace_prefix, ops
+SUPPORTED_HEAD_DIMS = tuple(range(8, 257, 8))
+COMPILED_HEAD_DIM_BUCKETS = (64, 96, 128, 256)
+SPLIT_HEAD_DIMS = SUPPORTED_HEAD_DIMS
 @dataclass(frozen=True)
     if any(int(v) <= 0 for v in values):
         raise ValueError("all shape values and num_sms must be positive")
     if int(head_dim) not in SUPPORTED_HEAD_DIMS:
+        raise ValueError("head_dim must be a positive multiple of 8 at most 256")
     block_n = 256 if head_dim <= 64 else (128 if head_dim <= 128 else 64)
     n_blocks = _ceildiv(seqlen_k, block_n)
     m_blocks = _ceildiv(seqlen_q, 64)
     if q.ndim != 4 or k.ndim != 4:
         raise ValueError("q and k must have shape (B, S, H, D)")
+    if q.shape[-1] not in SUPPORTED_HEAD_DIMS:
+        raise ValueError("head_dim must be a positive multiple of 8 at most 256")
     if num_sms is None:
         num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count
     splits = recommended_num_splits(
         device=q.device,
         dtype=torch.float32,
     )
+    d_rounded = (q.shape[3] + 31) & ~31
     out = torch.empty(
+        (splits, q.shape[0], q.shape[2], q.shape[1], d_rounded),
         device=q.device,
         dtype=torch.float32,
     )
 __all__ = [
     "FA2Workspace",
+    "COMPILED_HEAD_DIM_BUCKETS",
     "SPLIT_HEAD_DIMS",
     "SUPPORTED_HEAD_DIMS",
     "allocate_outputs",

build/torch212-cxx11-cu130-x86_64-linux/{_fa2_seqused_runtime_cuda_61bef7e.abi3.so → _fa2_seqused_runtime_cuda_99d26a1.abi3.so} RENAMED Viewed

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

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

build/torch212-cxx11-cu130-x86_64-linux/_ops.py CHANGED Viewed

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fa2-seqused-runtime",
-  "id": "_fa2_seqused_runtime_cuda_61bef7e",
   "version": 1,
   "license": "BSD-3-Clause",
   "python-depends": [],
@@ -16,9 +16,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
-      "_fa2_seqused_runtime_cuda_61bef7e.abi3.so": "Ma3gHi052hcrhMQ1yTUFX/7SzyXs+TeO5G5qMSk1Ch8=",
-      "_ops.py": "S5eXSORqSBWo6D7FltYgk8ttuDlbt6lSeHd0GmU8q6U=",
       "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   }

 {
   "name": "fa2-seqused-runtime",
+  "id": "_fa2_seqused_runtime_cuda_99d26a1",
   "version": 1,
   "license": "BSD-3-Clause",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "zaIQDw3dhB5xbCQPrZ9tGWr/Z559+ooyZOaGIa2/+LQ=",
+      "_fa2_seqused_runtime_cuda_99d26a1.abi3.so": "4dvtiwfQ7h1aRcHrps92BLHSNKQZ905NclVOCUViJOY=",
+      "_ops.py": "3zmyZsIXRYN9PWXNasLSKultcD9S0IaBbtFEFMbUth0=",
       "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   }