Kernels:

kernels-community
/

sage-attention

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on May 19

Commit

d0a1129

verified ·

1 Parent(s): 1042e4d

Uploaded using `kernel-builder`.

Browse files

Files changed (26) hide show

.gitattributes +6 -0
benchmarks/benchmark.py +82 -0
build/torch210-cxx11-cu128-aarch64-linux/_ops.py +3 -3
build/torch210-cxx11-cu128-aarch64-linux/{_sage_attention_cuda_a0cb380.abi3.so → _sage_attention_cuda_7535f60.abi3.so} +1 -1
build/torch210-cxx11-cu128-aarch64-linux/metadata.json +1 -1
build/torch210-cxx11-cu128-aarch64-linux/sm100_compile.py +7 -67
build/torch210-cxx11-cu130-aarch64-linux/_ops.py +3 -3
build/torch210-cxx11-cu130-aarch64-linux/{_sage_attention_cuda_a0cb380.abi3.so → _sage_attention_cuda_7535f60.abi3.so} +1 -1
build/torch210-cxx11-cu130-aarch64-linux/metadata.json +1 -1
build/torch210-cxx11-cu130-aarch64-linux/sm100_compile.py +7 -67
build/torch211-cxx11-cu128-aarch64-linux/_ops.py +3 -3
build/torch211-cxx11-cu128-aarch64-linux/{_sage_attention_cuda_a0cb380.abi3.so → _sage_attention_cuda_7535f60.abi3.so} +1 -1
build/torch211-cxx11-cu128-aarch64-linux/metadata.json +1 -1
build/torch211-cxx11-cu128-aarch64-linux/sm100_compile.py +7 -67
build/torch211-cxx11-cu130-aarch64-linux/_ops.py +3 -3
build/torch211-cxx11-cu130-aarch64-linux/{_sage_attention_cuda_a0cb380.abi3.so → _sage_attention_cuda_7535f60.abi3.so} +1 -1
build/torch211-cxx11-cu130-aarch64-linux/metadata.json +1 -1
build/torch211-cxx11-cu130-aarch64-linux/sm100_compile.py +7 -67
build/torch212-cxx11-cu130-aarch64-linux/_ops.py +3 -3
build/torch212-cxx11-cu130-aarch64-linux/{_sage_attention_cuda_a0cb380.abi3.so → _sage_attention_cuda_7535f60.abi3.so} +1 -1
build/torch212-cxx11-cu130-aarch64-linux/metadata.json +1 -1
build/torch212-cxx11-cu130-aarch64-linux/sm100_compile.py +7 -67
build/torch212-cxx11-cu132-aarch64-linux/_ops.py +3 -3
build/torch212-cxx11-cu132-aarch64-linux/{_sage_attention_cuda_a0cb380.abi3.so → _sage_attention_cuda_7535f60.abi3.so} +1 -1
build/torch212-cxx11-cu132-aarch64-linux/metadata.json +1 -1
build/torch212-cxx11-cu132-aarch64-linux/sm100_compile.py +7 -67

.gitattributes CHANGED Viewed

@@ -145,3 +145,9 @@ build/torch211-cxx11-cu128-aarch64-linux/_sage_attention_cuda_a0cb380.abi3.so fi
 build/torch211-cxx11-cu130-aarch64-linux/_sage_attention_cuda_a0cb380.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch212-cxx11-cu130-aarch64-linux/_sage_attention_cuda_a0cb380.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch212-cxx11-cu132-aarch64-linux/_sage_attention_cuda_a0cb380.abi3.so filter=lfs diff=lfs merge=lfs -text

 build/torch211-cxx11-cu130-aarch64-linux/_sage_attention_cuda_a0cb380.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch212-cxx11-cu130-aarch64-linux/_sage_attention_cuda_a0cb380.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch212-cxx11-cu132-aarch64-linux/_sage_attention_cuda_a0cb380.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu128-aarch64-linux/_sage_attention_cuda_7535f60.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu130-aarch64-linux/_sage_attention_cuda_7535f60.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch211-cxx11-cu128-aarch64-linux/_sage_attention_cuda_7535f60.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch211-cxx11-cu130-aarch64-linux/_sage_attention_cuda_7535f60.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch212-cxx11-cu130-aarch64-linux/_sage_attention_cuda_7535f60.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch212-cxx11-cu132-aarch64-linux/_sage_attention_cuda_7535f60.abi3.so filter=lfs diff=lfs merge=lfs -text

benchmarks/benchmark.py ADDED Viewed

	@@ -0,0 +1,82 @@

+import torch
+import torch.nn.functional as F
+from kernels.benchmark import Benchmark
+# SageAttention is approximate (INT8 quantized QK) so element-wise allclose
+# is too strict.  Use cosine similarity instead (threshold 0.99).
+_orig_allclose = torch.allclose
+torch.allclose = lambda a, b, **_kw: (
+    F.cosine_similarity(a.flatten().float().unsqueeze(0),
+                        b.flatten().float().unsqueeze(0)).item() > 0.99
+)
+def _ref(q, k, v, is_causal=False):
+    return F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
+class SageAttentionBenchmark(Benchmark):
+    seed: int = 42
+    # --- base: B=2, H=32, L=1024, D=128 ---
+    def setup_base(self):
+        B, H, L, D = 2, 32, 1024, 128
+        self.q = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=self.device)
+        self.k = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=self.device)
+        self.v = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=self.device)
+        self.out = torch.empty_like(self.q)
+    def benchmark_base(self):
+        self.out = self.kernel.sageattn(self.q, self.k, self.v, tensor_layout="HND")
+    def verify_base(self) -> torch.Tensor:
+        return _ref(self.q, self.k, self.v)
+    # --- causal: B=2, H=32, L=1024, D=128 with causal mask ---
+    def setup_causal(self):
+        B, H, L, D = 2, 32, 1024, 128
+        self.q = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=self.device)
+        self.k = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=self.device)
+        self.v = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=self.device)
+        self.out = torch.empty_like(self.q)
+    def benchmark_causal(self):
+        self.out = self.kernel.sageattn(
+            self.q, self.k, self.v, tensor_layout="HND", is_causal=True
+        )
+    def verify_causal(self) -> torch.Tensor:
+        return _ref(self.q, self.k, self.v, is_causal=True)
+    # --- large: B=4, H=32, L=4096, D=128 ---
+    def setup_large(self):
+        B, H, L, D = 4, 32, 4096, 128
+        self.q = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=self.device)
+        self.k = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=self.device)
+        self.v = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=self.device)
+        self.out = torch.empty_like(self.q)
+    def benchmark_large(self):
+        self.out = self.kernel.sageattn(self.q, self.k, self.v, tensor_layout="HND")
+    def verify_large(self) -> torch.Tensor:
+        return _ref(self.q, self.k, self.v)
+    # --- d64: B=4, H=32, L=2048, D=64 (smaller head dim) ---
+    def setup_d64(self):
+        B, H, L, D = 4, 32, 2048, 64
+        self.q = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=self.device)
+        self.k = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=self.device)
+        self.v = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=self.device)
+        self.out = torch.empty_like(self.q)
+    def benchmark_d64(self):
+        self.out = self.kernel.sageattn(self.q, self.k, self.v, tensor_layout="HND")
+    def verify_d64(self) -> torch.Tensor:
+        return _ref(self.q, self.k, self.v)

build/torch210-cxx11-cu128-aarch64-linux/_ops.py CHANGED Viewed

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

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

build/torch210-cxx11-cu128-aarch64-linux/{_sage_attention_cuda_a0cb380.abi3.so → _sage_attention_cuda_7535f60.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:256efe8e854ed0d7e4e200b7ba56adf26b8361997eb461635cad981708ef694e
 size 33330136

build/torch210-cxx11-cu128-aarch64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "sage-attention",
-  "id": "_sage_attention_cuda_a0cb380",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "sage-attention",
+  "id": "_sage_attention_cuda_7535f60",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch210-cxx11-cu128-aarch64-linux/sm100_compile.py CHANGED Viewed

@@ -28,31 +28,6 @@ from torch.nn.functional import scaled_dot_product_attention as sdpa
 # Low-level ops with torch.compile support (custom_op + register_fake)
 # ---------------------------------------------------------------------------
-@torch.library.custom_op(
-    add_op_namespace_prefix("mha_fwd"), mutates_args=(), device_types="cuda"
-)
-def mha_fwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    sfq: torch.Tensor,
-    sfk: torch.Tensor,
-    sfv: torch.Tensor,
-    delta_s: torch.Tensor,
-    unpadded_k: int,
-    out: Optional[torch.Tensor],
-    softmax_scale: float,
-    is_causal: bool,
-    per_block_mean: bool,
-    is_bf16: bool,
-) -> List[torch.Tensor]:
-    return ops.mha_fwd(
-        q, k, v, sfq, sfk, sfv, delta_s,
-        unpadded_k, out, softmax_scale, is_causal,
-        per_block_mean, is_bf16,
-    )
 @torch.library.register_fake(add_op_namespace_prefix("mha_fwd"))
 def mha_fwd_fake(
     q: torch.Tensor,
@@ -86,20 +61,6 @@ def mha_fwd_fake(
     return [fake_out, fake_lse]
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant"))
 def scaled_fp4_quant_fake(
     input: torch.Tensor,
@@ -110,20 +71,6 @@ def scaled_fp4_quant_fake(
     pass
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant_permute"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant_permute(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant_permute(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_permute"))
 def scaled_fp4_quant_permute_fake(
     input: torch.Tensor,
@@ -134,20 +81,6 @@ def scaled_fp4_quant_permute_fake(
     pass
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant_trans"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant_trans(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant_trans(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_trans"))
 def scaled_fp4_quant_trans_fake(
     input: torch.Tensor,
@@ -158,6 +91,13 @@ def scaled_fp4_quant_trans_fake(
     pass
 # ---------------------------------------------------------------------------
 # Triton kernel for grouped mean subtraction
 # ---------------------------------------------------------------------------

 # Low-level ops with torch.compile support (custom_op + register_fake)
 # ---------------------------------------------------------------------------
 @torch.library.register_fake(add_op_namespace_prefix("mha_fwd"))
 def mha_fwd_fake(
     q: torch.Tensor,
     return [fake_out, fake_lse]
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant"))
 def scaled_fp4_quant_fake(
     input: torch.Tensor,
     pass
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_permute"))
 def scaled_fp4_quant_permute_fake(
     input: torch.Tensor,
     pass
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_trans"))
 def scaled_fp4_quant_trans_fake(
     input: torch.Tensor,
     pass
+# Direct references to C++ ops (same pattern as sm89_compile.py)
+mha_fwd = ops.mha_fwd
+scaled_fp4_quant = ops.scaled_fp4_quant
+scaled_fp4_quant_permute = ops.scaled_fp4_quant_permute
+scaled_fp4_quant_trans = ops.scaled_fp4_quant_trans
 # ---------------------------------------------------------------------------
 # Triton kernel for grouped mean subtraction
 # ---------------------------------------------------------------------------

build/torch210-cxx11-cu130-aarch64-linux/_ops.py CHANGED Viewed

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

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

build/torch210-cxx11-cu130-aarch64-linux/{_sage_attention_cuda_a0cb380.abi3.so → _sage_attention_cuda_7535f60.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:810539bd3e29dd2892d1b3fe396d7e5ad676cfdf524b884e7f9e1b6483a6369a
 size 33810512

build/torch210-cxx11-cu130-aarch64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "sage-attention",
-  "id": "_sage_attention_cuda_a0cb380",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "sage-attention",
+  "id": "_sage_attention_cuda_7535f60",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch210-cxx11-cu130-aarch64-linux/sm100_compile.py CHANGED Viewed

@@ -28,31 +28,6 @@ from torch.nn.functional import scaled_dot_product_attention as sdpa
 # Low-level ops with torch.compile support (custom_op + register_fake)
 # ---------------------------------------------------------------------------
-@torch.library.custom_op(
-    add_op_namespace_prefix("mha_fwd"), mutates_args=(), device_types="cuda"
-)
-def mha_fwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    sfq: torch.Tensor,
-    sfk: torch.Tensor,
-    sfv: torch.Tensor,
-    delta_s: torch.Tensor,
-    unpadded_k: int,
-    out: Optional[torch.Tensor],
-    softmax_scale: float,
-    is_causal: bool,
-    per_block_mean: bool,
-    is_bf16: bool,
-) -> List[torch.Tensor]:
-    return ops.mha_fwd(
-        q, k, v, sfq, sfk, sfv, delta_s,
-        unpadded_k, out, softmax_scale, is_causal,
-        per_block_mean, is_bf16,
-    )
 @torch.library.register_fake(add_op_namespace_prefix("mha_fwd"))
 def mha_fwd_fake(
     q: torch.Tensor,
@@ -86,20 +61,6 @@ def mha_fwd_fake(
     return [fake_out, fake_lse]
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant"))
 def scaled_fp4_quant_fake(
     input: torch.Tensor,
@@ -110,20 +71,6 @@ def scaled_fp4_quant_fake(
     pass
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant_permute"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant_permute(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant_permute(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_permute"))
 def scaled_fp4_quant_permute_fake(
     input: torch.Tensor,
@@ -134,20 +81,6 @@ def scaled_fp4_quant_permute_fake(
     pass
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant_trans"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant_trans(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant_trans(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_trans"))
 def scaled_fp4_quant_trans_fake(
     input: torch.Tensor,
@@ -158,6 +91,13 @@ def scaled_fp4_quant_trans_fake(
     pass
 # ---------------------------------------------------------------------------
 # Triton kernel for grouped mean subtraction
 # ---------------------------------------------------------------------------

 # Low-level ops with torch.compile support (custom_op + register_fake)
 # ---------------------------------------------------------------------------
 @torch.library.register_fake(add_op_namespace_prefix("mha_fwd"))
 def mha_fwd_fake(
     q: torch.Tensor,
     return [fake_out, fake_lse]
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant"))
 def scaled_fp4_quant_fake(
     input: torch.Tensor,
     pass
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_permute"))
 def scaled_fp4_quant_permute_fake(
     input: torch.Tensor,
     pass
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_trans"))
 def scaled_fp4_quant_trans_fake(
     input: torch.Tensor,
     pass
+# Direct references to C++ ops (same pattern as sm89_compile.py)
+mha_fwd = ops.mha_fwd
+scaled_fp4_quant = ops.scaled_fp4_quant
+scaled_fp4_quant_permute = ops.scaled_fp4_quant_permute
+scaled_fp4_quant_trans = ops.scaled_fp4_quant_trans
 # ---------------------------------------------------------------------------
 # Triton kernel for grouped mean subtraction
 # ---------------------------------------------------------------------------

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

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

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

build/torch211-cxx11-cu128-aarch64-linux/{_sage_attention_cuda_a0cb380.abi3.so → _sage_attention_cuda_7535f60.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:10c680f7b61513fc2eea02bdb61b9838776215d61c7fb50a732c3d493c10e695
 size 33326264

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

@@ -1,6 +1,6 @@
 {
   "name": "sage-attention",
-  "id": "_sage_attention_cuda_a0cb380",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "sage-attention",
+  "id": "_sage_attention_cuda_7535f60",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch211-cxx11-cu128-aarch64-linux/sm100_compile.py CHANGED Viewed

@@ -28,31 +28,6 @@ from torch.nn.functional import scaled_dot_product_attention as sdpa
 # Low-level ops with torch.compile support (custom_op + register_fake)
 # ---------------------------------------------------------------------------
-@torch.library.custom_op(
-    add_op_namespace_prefix("mha_fwd"), mutates_args=(), device_types="cuda"
-)
-def mha_fwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    sfq: torch.Tensor,
-    sfk: torch.Tensor,
-    sfv: torch.Tensor,
-    delta_s: torch.Tensor,
-    unpadded_k: int,
-    out: Optional[torch.Tensor],
-    softmax_scale: float,
-    is_causal: bool,
-    per_block_mean: bool,
-    is_bf16: bool,
-) -> List[torch.Tensor]:
-    return ops.mha_fwd(
-        q, k, v, sfq, sfk, sfv, delta_s,
-        unpadded_k, out, softmax_scale, is_causal,
-        per_block_mean, is_bf16,
-    )
 @torch.library.register_fake(add_op_namespace_prefix("mha_fwd"))
 def mha_fwd_fake(
     q: torch.Tensor,
@@ -86,20 +61,6 @@ def mha_fwd_fake(
     return [fake_out, fake_lse]
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant"))
 def scaled_fp4_quant_fake(
     input: torch.Tensor,
@@ -110,20 +71,6 @@ def scaled_fp4_quant_fake(
     pass
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant_permute"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant_permute(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant_permute(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_permute"))
 def scaled_fp4_quant_permute_fake(
     input: torch.Tensor,
@@ -134,20 +81,6 @@ def scaled_fp4_quant_permute_fake(
     pass
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant_trans"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant_trans(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant_trans(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_trans"))
 def scaled_fp4_quant_trans_fake(
     input: torch.Tensor,
@@ -158,6 +91,13 @@ def scaled_fp4_quant_trans_fake(
     pass
 # ---------------------------------------------------------------------------
 # Triton kernel for grouped mean subtraction
 # ---------------------------------------------------------------------------

 # Low-level ops with torch.compile support (custom_op + register_fake)
 # ---------------------------------------------------------------------------
 @torch.library.register_fake(add_op_namespace_prefix("mha_fwd"))
 def mha_fwd_fake(
     q: torch.Tensor,
     return [fake_out, fake_lse]
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant"))
 def scaled_fp4_quant_fake(
     input: torch.Tensor,
     pass
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_permute"))
 def scaled_fp4_quant_permute_fake(
     input: torch.Tensor,
     pass
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_trans"))
 def scaled_fp4_quant_trans_fake(
     input: torch.Tensor,
     pass
+# Direct references to C++ ops (same pattern as sm89_compile.py)
+mha_fwd = ops.mha_fwd
+scaled_fp4_quant = ops.scaled_fp4_quant
+scaled_fp4_quant_permute = ops.scaled_fp4_quant_permute
+scaled_fp4_quant_trans = ops.scaled_fp4_quant_trans
 # ---------------------------------------------------------------------------
 # Triton kernel for grouped mean subtraction
 # ---------------------------------------------------------------------------

build/torch211-cxx11-cu130-aarch64-linux/_ops.py CHANGED Viewed

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

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

build/torch211-cxx11-cu130-aarch64-linux/{_sage_attention_cuda_a0cb380.abi3.so → _sage_attention_cuda_7535f60.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:51ec2a4ca4bb8f8002a971338920422bd21471a2bca43f63712a427c50b21cd2
 size 33806648

build/torch211-cxx11-cu130-aarch64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "sage-attention",
-  "id": "_sage_attention_cuda_a0cb380",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "sage-attention",
+  "id": "_sage_attention_cuda_7535f60",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch211-cxx11-cu130-aarch64-linux/sm100_compile.py CHANGED Viewed

@@ -28,31 +28,6 @@ from torch.nn.functional import scaled_dot_product_attention as sdpa
 # Low-level ops with torch.compile support (custom_op + register_fake)
 # ---------------------------------------------------------------------------
-@torch.library.custom_op(
-    add_op_namespace_prefix("mha_fwd"), mutates_args=(), device_types="cuda"
-)
-def mha_fwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    sfq: torch.Tensor,
-    sfk: torch.Tensor,
-    sfv: torch.Tensor,
-    delta_s: torch.Tensor,
-    unpadded_k: int,
-    out: Optional[torch.Tensor],
-    softmax_scale: float,
-    is_causal: bool,
-    per_block_mean: bool,
-    is_bf16: bool,
-) -> List[torch.Tensor]:
-    return ops.mha_fwd(
-        q, k, v, sfq, sfk, sfv, delta_s,
-        unpadded_k, out, softmax_scale, is_causal,
-        per_block_mean, is_bf16,
-    )
 @torch.library.register_fake(add_op_namespace_prefix("mha_fwd"))
 def mha_fwd_fake(
     q: torch.Tensor,
@@ -86,20 +61,6 @@ def mha_fwd_fake(
     return [fake_out, fake_lse]
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant"))
 def scaled_fp4_quant_fake(
     input: torch.Tensor,
@@ -110,20 +71,6 @@ def scaled_fp4_quant_fake(
     pass
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant_permute"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant_permute(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant_permute(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_permute"))
 def scaled_fp4_quant_permute_fake(
     input: torch.Tensor,
@@ -134,20 +81,6 @@ def scaled_fp4_quant_permute_fake(
     pass
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant_trans"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant_trans(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant_trans(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_trans"))
 def scaled_fp4_quant_trans_fake(
     input: torch.Tensor,
@@ -158,6 +91,13 @@ def scaled_fp4_quant_trans_fake(
     pass
 # ---------------------------------------------------------------------------
 # Triton kernel for grouped mean subtraction
 # ---------------------------------------------------------------------------

 # Low-level ops with torch.compile support (custom_op + register_fake)
 # ---------------------------------------------------------------------------
 @torch.library.register_fake(add_op_namespace_prefix("mha_fwd"))
 def mha_fwd_fake(
     q: torch.Tensor,
     return [fake_out, fake_lse]
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant"))
 def scaled_fp4_quant_fake(
     input: torch.Tensor,
     pass
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_permute"))
 def scaled_fp4_quant_permute_fake(
     input: torch.Tensor,
     pass
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_trans"))
 def scaled_fp4_quant_trans_fake(
     input: torch.Tensor,
     pass
+# Direct references to C++ ops (same pattern as sm89_compile.py)
+mha_fwd = ops.mha_fwd
+scaled_fp4_quant = ops.scaled_fp4_quant
+scaled_fp4_quant_permute = ops.scaled_fp4_quant_permute
+scaled_fp4_quant_trans = ops.scaled_fp4_quant_trans
 # ---------------------------------------------------------------------------
 # Triton kernel for grouped mean subtraction
 # ---------------------------------------------------------------------------

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

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

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

build/torch212-cxx11-cu130-aarch64-linux/{_sage_attention_cuda_a0cb380.abi3.so → _sage_attention_cuda_7535f60.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:3d4a3a414a1677470451d1dde012fbc31666990c8cb6701843f9b1bccd0a05d4
 size 33810192

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

@@ -1,6 +1,6 @@
 {
   "name": "sage-attention",
-  "id": "_sage_attention_cuda_a0cb380",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "sage-attention",
+  "id": "_sage_attention_cuda_7535f60",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch212-cxx11-cu130-aarch64-linux/sm100_compile.py CHANGED Viewed

@@ -28,31 +28,6 @@ from torch.nn.functional import scaled_dot_product_attention as sdpa
 # Low-level ops with torch.compile support (custom_op + register_fake)
 # ---------------------------------------------------------------------------
-@torch.library.custom_op(
-    add_op_namespace_prefix("mha_fwd"), mutates_args=(), device_types="cuda"
-)
-def mha_fwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    sfq: torch.Tensor,
-    sfk: torch.Tensor,
-    sfv: torch.Tensor,
-    delta_s: torch.Tensor,
-    unpadded_k: int,
-    out: Optional[torch.Tensor],
-    softmax_scale: float,
-    is_causal: bool,
-    per_block_mean: bool,
-    is_bf16: bool,
-) -> List[torch.Tensor]:
-    return ops.mha_fwd(
-        q, k, v, sfq, sfk, sfv, delta_s,
-        unpadded_k, out, softmax_scale, is_causal,
-        per_block_mean, is_bf16,
-    )
 @torch.library.register_fake(add_op_namespace_prefix("mha_fwd"))
 def mha_fwd_fake(
     q: torch.Tensor,
@@ -86,20 +61,6 @@ def mha_fwd_fake(
     return [fake_out, fake_lse]
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant"))
 def scaled_fp4_quant_fake(
     input: torch.Tensor,
@@ -110,20 +71,6 @@ def scaled_fp4_quant_fake(
     pass
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant_permute"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant_permute(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant_permute(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_permute"))
 def scaled_fp4_quant_permute_fake(
     input: torch.Tensor,
@@ -134,20 +81,6 @@ def scaled_fp4_quant_permute_fake(
     pass
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant_trans"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant_trans(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant_trans(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_trans"))
 def scaled_fp4_quant_trans_fake(
     input: torch.Tensor,
@@ -158,6 +91,13 @@ def scaled_fp4_quant_trans_fake(
     pass
 # ---------------------------------------------------------------------------
 # Triton kernel for grouped mean subtraction
 # ---------------------------------------------------------------------------

 # Low-level ops with torch.compile support (custom_op + register_fake)
 # ---------------------------------------------------------------------------
 @torch.library.register_fake(add_op_namespace_prefix("mha_fwd"))
 def mha_fwd_fake(
     q: torch.Tensor,
     return [fake_out, fake_lse]
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant"))
 def scaled_fp4_quant_fake(
     input: torch.Tensor,
     pass
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_permute"))
 def scaled_fp4_quant_permute_fake(
     input: torch.Tensor,
     pass
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_trans"))
 def scaled_fp4_quant_trans_fake(
     input: torch.Tensor,
     pass
+# Direct references to C++ ops (same pattern as sm89_compile.py)
+mha_fwd = ops.mha_fwd
+scaled_fp4_quant = ops.scaled_fp4_quant
+scaled_fp4_quant_permute = ops.scaled_fp4_quant_permute
+scaled_fp4_quant_trans = ops.scaled_fp4_quant_trans
 # ---------------------------------------------------------------------------
 # Triton kernel for grouped mean subtraction
 # ---------------------------------------------------------------------------

build/torch212-cxx11-cu132-aarch64-linux/_ops.py CHANGED Viewed

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

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

build/torch212-cxx11-cu132-aarch64-linux/{_sage_attention_cuda_a0cb380.abi3.so → _sage_attention_cuda_7535f60.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:7a34459318d6736e53202e2d7244d1a609d533273ca31f5cf0c3184dcfdf1a05
 size 33753872

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

@@ -1,6 +1,6 @@
 {
   "name": "sage-attention",
-  "id": "_sage_attention_cuda_a0cb380",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "sage-attention",
+  "id": "_sage_attention_cuda_7535f60",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch212-cxx11-cu132-aarch64-linux/sm100_compile.py CHANGED Viewed

@@ -28,31 +28,6 @@ from torch.nn.functional import scaled_dot_product_attention as sdpa
 # Low-level ops with torch.compile support (custom_op + register_fake)
 # ---------------------------------------------------------------------------
-@torch.library.custom_op(
-    add_op_namespace_prefix("mha_fwd"), mutates_args=(), device_types="cuda"
-)
-def mha_fwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    sfq: torch.Tensor,
-    sfk: torch.Tensor,
-    sfv: torch.Tensor,
-    delta_s: torch.Tensor,
-    unpadded_k: int,
-    out: Optional[torch.Tensor],
-    softmax_scale: float,
-    is_causal: bool,
-    per_block_mean: bool,
-    is_bf16: bool,
-) -> List[torch.Tensor]:
-    return ops.mha_fwd(
-        q, k, v, sfq, sfk, sfv, delta_s,
-        unpadded_k, out, softmax_scale, is_causal,
-        per_block_mean, is_bf16,
-    )
 @torch.library.register_fake(add_op_namespace_prefix("mha_fwd"))
 def mha_fwd_fake(
     q: torch.Tensor,
@@ -86,20 +61,6 @@ def mha_fwd_fake(
     return [fake_out, fake_lse]
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant"))
 def scaled_fp4_quant_fake(
     input: torch.Tensor,
@@ -110,20 +71,6 @@ def scaled_fp4_quant_fake(
     pass
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant_permute"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant_permute(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant_permute(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_permute"))
 def scaled_fp4_quant_permute_fake(
     input: torch.Tensor,
@@ -134,20 +81,6 @@ def scaled_fp4_quant_permute_fake(
     pass
-@torch.library.custom_op(
-    add_op_namespace_prefix("scaled_fp4_quant_trans"),
-    mutates_args=("output", "output_sf"),
-    device_types="cuda",
-)
-def scaled_fp4_quant_trans(
-    input: torch.Tensor,
-    output: torch.Tensor,
-    output_sf: torch.Tensor,
-    tensor_layout: int,
-) -> None:
-    ops.scaled_fp4_quant_trans(input, output, output_sf, tensor_layout)
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_trans"))
 def scaled_fp4_quant_trans_fake(
     input: torch.Tensor,
@@ -158,6 +91,13 @@ def scaled_fp4_quant_trans_fake(
     pass
 # ---------------------------------------------------------------------------
 # Triton kernel for grouped mean subtraction
 # ---------------------------------------------------------------------------

 # Low-level ops with torch.compile support (custom_op + register_fake)
 # ---------------------------------------------------------------------------
 @torch.library.register_fake(add_op_namespace_prefix("mha_fwd"))
 def mha_fwd_fake(
     q: torch.Tensor,
     return [fake_out, fake_lse]
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant"))
 def scaled_fp4_quant_fake(
     input: torch.Tensor,
     pass
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_permute"))
 def scaled_fp4_quant_permute_fake(
     input: torch.Tensor,
     pass
 @torch.library.register_fake(add_op_namespace_prefix("scaled_fp4_quant_trans"))
 def scaled_fp4_quant_trans_fake(
     input: torch.Tensor,
     pass
+# Direct references to C++ ops (same pattern as sm89_compile.py)
+mha_fwd = ops.mha_fwd
+scaled_fp4_quant = ops.scaled_fp4_quant
+scaled_fp4_quant_permute = ops.scaled_fp4_quant_permute
+scaled_fp4_quant_trans = ops.scaled_fp4_quant_trans
 # ---------------------------------------------------------------------------
 # Triton kernel for grouped mean subtraction
 # ---------------------------------------------------------------------------