Build uploaded using `kernels`.

.gitattributes

@@ -55,3 +55,9 @@ build/torch210-cxx11-cu130-aarch64-linux/_sage_attention_cuda_4eabbf5.abi3.so fi
 build/torch29-cxx11-cu126-aarch64-linux/_sage_attention_cuda_4eabbf5.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch29-cxx11-cu128-aarch64-linux/_sage_attention_cuda_4eabbf5.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch29-cxx11-cu130-aarch64-linux/_sage_attention_cuda_4eabbf5.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu126-x86_64-linux/_sage_attention_cuda_4eabbf5.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu128-x86_64-linux/_sage_attention_cuda_4eabbf5.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu130-x86_64-linux/_sage_attention_cuda_4eabbf5.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-cxx11-cu126-x86_64-linux/_sage_attention_cuda_4eabbf5.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-cxx11-cu128-x86_64-linux/_sage_attention_cuda_4eabbf5.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-cxx11-cu130-x86_64-linux/_sage_attention_cuda_4eabbf5.abi3.so filter=lfs diff=lfs merge=lfs -text

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

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

build/torch210-cxx11-cu126-x86_64-linux/{_sage_attention_6e51d70.abi3.so → _sage_attention_cuda_4eabbf5.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8b4b01fb9797a2b69820d313342d84dcfc94f0bf333aec597283953b8bd9eb7f
-size 26060392
+oid sha256:38b7d44512d6188693316247322f279ef483bc2801118392dae59e12578a471d
+size 26060416

build/torch210-cxx11-cu126-x86_64-linux/core.py

@@ -26,21 +26,33 @@ from .quant import per_channel_fp8
 from .quant_per_thread import per_thread_int8 as per_thread_int8_triton
 
 try:
-    from . import sm80_compile
+    from .sm80_compile import (
+        qk_int8_sv_f16_accum_f32_attn as sm80_qk_int8_sv_f16_accum_f32_attn,
+        qk_int8_sv_f16_accum_f16_fuse_v_mean_attn as sm80_qk_int8_sv_f16_accum_f16_fuse_v_mean_attn,
+        qk_int8_sv_f16_accum_f16_attn as sm80_qk_int8_sv_f16_accum_f16_attn,
+        qk_int8_sv_f16_accum_f16_attn_inst_buf as sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf,
+    )
     SM80_ENABLED = True
-except:
+except Exception:
     SM80_ENABLED = False
 
 try:
-    from . import sm89_compile
+    from .sm89_compile import (
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn,
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn,
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf,
+        qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf,
+    )
     SM89_ENABLED = True
-except:
+except Exception:
     SM89_ENABLED = False
 
 try:
-    from . import sm90_compile
+    from .sm90_compile import (
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90 as sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90,
+    )
     SM90_ENABLED = True
-except:
+except Exception:
     SM90_ENABLED = False
 
 from typing import Any, List, Literal, Optional, Tuple, Union
@@ -378,7 +390,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
 
     if pv_accum_dtype == "fp32":
         v = v.to(torch.float16)
-        lse = sm80_compile.qk_int8_sv_f16_accum_f32_attn(
+        lse = sm80_qk_int8_sv_f16_accum_f32_attn(
             q_int8,
             k_int8,
             v,
@@ -394,7 +406,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
     elif pv_accum_dtype == "fp16":
         if smooth_v:
             smoothed_v, vm = sub_mean(v, tensor_layout=tensor_layout)
-            lse = sm80_compile.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+            lse = sm80_qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 smoothed_v,
@@ -410,7 +422,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
         else:
             v = v.to(torch.float16)
-            lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn(
+            lse = sm80_qk_int8_sv_f16_accum_f16_attn(
                 q_int8,
                 k_int8,
                 v,
@@ -425,7 +437,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
     elif pv_accum_dtype == "fp16+fp32":
         v = v.to(torch.float16)
-        lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        lse = sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf(
             q_int8,
             k_int8,
             v,
@@ -650,7 +662,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
     )
     if pv_accum_dtype == "fp32":
         if smooth_v:
-            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+            lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -667,7 +679,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
         else:
-            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+            lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -683,7 +695,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp32":
-        lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -699,7 +711,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp16":
-        lse = sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        lse = sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -936,7 +948,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
 
     if pv_accum_dtype == "fp32":
         raise NotImplementedError("Please use pv_accum_dtype='fp32+fp32' for sm90.")
-        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        lse = sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
             q_int8,
             k_int8,
             v_fp8,
@@ -951,7 +963,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
             _return_lse,
         )
     elif pv_accum_dtype == "fp32+fp32":
-        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        lse = sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
             q_int8,
             k_int8,
             v_fp8,

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

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

build/torch210-cxx11-cu128-x86_64-linux/{_sage_attention_6e51d70.abi3.so → _sage_attention_cuda_4eabbf5.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:86cee1d97e0e6fcd5aaf9480aaa092d96d7e8780214e2768ae1a68448df2dbf3
-size 26573256
+oid sha256:7fedc1235631ddd8ed0db9f8fd9e7860a9735acbe007f0d3cd81d4ecc8a5a0c6
+size 26573280

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

@@ -26,21 +26,33 @@ from .quant import per_channel_fp8
 from .quant_per_thread import per_thread_int8 as per_thread_int8_triton
 
 try:
-    from . import sm80_compile
+    from .sm80_compile import (
+        qk_int8_sv_f16_accum_f32_attn as sm80_qk_int8_sv_f16_accum_f32_attn,
+        qk_int8_sv_f16_accum_f16_fuse_v_mean_attn as sm80_qk_int8_sv_f16_accum_f16_fuse_v_mean_attn,
+        qk_int8_sv_f16_accum_f16_attn as sm80_qk_int8_sv_f16_accum_f16_attn,
+        qk_int8_sv_f16_accum_f16_attn_inst_buf as sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf,
+    )
     SM80_ENABLED = True
-except:
+except Exception:
     SM80_ENABLED = False
 
 try:
-    from . import sm89_compile
+    from .sm89_compile import (
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn,
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn,
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf,
+        qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf,
+    )
     SM89_ENABLED = True
-except:
+except Exception:
     SM89_ENABLED = False
 
 try:
-    from . import sm90_compile
+    from .sm90_compile import (
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90 as sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90,
+    )
     SM90_ENABLED = True
-except:
+except Exception:
     SM90_ENABLED = False
 
 from typing import Any, List, Literal, Optional, Tuple, Union
@@ -378,7 +390,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
 
     if pv_accum_dtype == "fp32":
         v = v.to(torch.float16)
-        lse = sm80_compile.qk_int8_sv_f16_accum_f32_attn(
+        lse = sm80_qk_int8_sv_f16_accum_f32_attn(
             q_int8,
             k_int8,
             v,
@@ -394,7 +406,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
     elif pv_accum_dtype == "fp16":
         if smooth_v:
             smoothed_v, vm = sub_mean(v, tensor_layout=tensor_layout)
-            lse = sm80_compile.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+            lse = sm80_qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 smoothed_v,
@@ -410,7 +422,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
         else:
             v = v.to(torch.float16)
-            lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn(
+            lse = sm80_qk_int8_sv_f16_accum_f16_attn(
                 q_int8,
                 k_int8,
                 v,
@@ -425,7 +437,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
     elif pv_accum_dtype == "fp16+fp32":
         v = v.to(torch.float16)
-        lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        lse = sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf(
             q_int8,
             k_int8,
             v,
@@ -650,7 +662,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
     )
     if pv_accum_dtype == "fp32":
         if smooth_v:
-            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+            lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -667,7 +679,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
         else:
-            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+            lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -683,7 +695,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp32":
-        lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -699,7 +711,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp16":
-        lse = sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        lse = sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -936,7 +948,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
 
     if pv_accum_dtype == "fp32":
         raise NotImplementedError("Please use pv_accum_dtype='fp32+fp32' for sm90.")
-        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        lse = sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
             q_int8,
             k_int8,
             v_fp8,
@@ -951,7 +963,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
             _return_lse,
         )
     elif pv_accum_dtype == "fp32+fp32":
-        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        lse = sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
             q_int8,
             k_int8,
             v_fp8,

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

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

build/torch210-cxx11-cu130-x86_64-linux/{_sage_attention_6e51d70.abi3.so → _sage_attention_cuda_4eabbf5.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:52f21795a824fc986e20cc3bda1ccfb7327343868888692d117ee8b81376dda9
-size 27244536
+oid sha256:8a6aaec766a6189b2cf7d38059332368797a2bc3c8e3a8ceccd8620aa78ec9a8
+size 27244560

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

@@ -26,21 +26,33 @@ from .quant import per_channel_fp8
 from .quant_per_thread import per_thread_int8 as per_thread_int8_triton
 
 try:
-    from . import sm80_compile
+    from .sm80_compile import (
+        qk_int8_sv_f16_accum_f32_attn as sm80_qk_int8_sv_f16_accum_f32_attn,
+        qk_int8_sv_f16_accum_f16_fuse_v_mean_attn as sm80_qk_int8_sv_f16_accum_f16_fuse_v_mean_attn,
+        qk_int8_sv_f16_accum_f16_attn as sm80_qk_int8_sv_f16_accum_f16_attn,
+        qk_int8_sv_f16_accum_f16_attn_inst_buf as sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf,
+    )
     SM80_ENABLED = True
-except:
+except Exception:
     SM80_ENABLED = False
 
 try:
-    from . import sm89_compile
+    from .sm89_compile import (
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn,
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn,
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf,
+        qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf,
+    )
     SM89_ENABLED = True
-except:
+except Exception:
     SM89_ENABLED = False
 
 try:
-    from . import sm90_compile
+    from .sm90_compile import (
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90 as sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90,
+    )
     SM90_ENABLED = True
-except:
+except Exception:
     SM90_ENABLED = False
 
 from typing import Any, List, Literal, Optional, Tuple, Union
@@ -378,7 +390,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
 
     if pv_accum_dtype == "fp32":
         v = v.to(torch.float16)
-        lse = sm80_compile.qk_int8_sv_f16_accum_f32_attn(
+        lse = sm80_qk_int8_sv_f16_accum_f32_attn(
             q_int8,
             k_int8,
             v,
@@ -394,7 +406,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
     elif pv_accum_dtype == "fp16":
         if smooth_v:
             smoothed_v, vm = sub_mean(v, tensor_layout=tensor_layout)
-            lse = sm80_compile.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+            lse = sm80_qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 smoothed_v,
@@ -410,7 +422,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
         else:
             v = v.to(torch.float16)
-            lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn(
+            lse = sm80_qk_int8_sv_f16_accum_f16_attn(
                 q_int8,
                 k_int8,
                 v,
@@ -425,7 +437,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
     elif pv_accum_dtype == "fp16+fp32":
         v = v.to(torch.float16)
-        lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        lse = sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf(
             q_int8,
             k_int8,
             v,
@@ -650,7 +662,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
     )
     if pv_accum_dtype == "fp32":
         if smooth_v:
-            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+            lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -667,7 +679,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
         else:
-            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+            lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -683,7 +695,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp32":
-        lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -699,7 +711,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp16":
-        lse = sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        lse = sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -936,7 +948,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
 
     if pv_accum_dtype == "fp32":
         raise NotImplementedError("Please use pv_accum_dtype='fp32+fp32' for sm90.")
-        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        lse = sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
             q_int8,
             k_int8,
             v_fp8,
@@ -951,7 +963,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
             _return_lse,
         )
     elif pv_accum_dtype == "fp32+fp32":
-        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        lse = sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
             q_int8,
             k_int8,
             v_fp8,

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

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

build/torch29-cxx11-cu126-x86_64-linux/{_sage_attention_6e51d70.abi3.so → _sage_attention_cuda_4eabbf5.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b3aaf3cffb71067f4816ba5da358cd3a0d3b19b9cb2ba1ec0eedf1bf9b4051e3
-size 26050976
+oid sha256:ff22b0633165c491f8fbd9c8014a0534313c3c8b059fe6ffe3a13bfc09ba3ed1
+size 26051000

build/torch29-cxx11-cu126-x86_64-linux/core.py

@@ -26,21 +26,33 @@ from .quant import per_channel_fp8
 from .quant_per_thread import per_thread_int8 as per_thread_int8_triton
 
 try:
-    from . import sm80_compile
+    from .sm80_compile import (
+        qk_int8_sv_f16_accum_f32_attn as sm80_qk_int8_sv_f16_accum_f32_attn,
+        qk_int8_sv_f16_accum_f16_fuse_v_mean_attn as sm80_qk_int8_sv_f16_accum_f16_fuse_v_mean_attn,
+        qk_int8_sv_f16_accum_f16_attn as sm80_qk_int8_sv_f16_accum_f16_attn,
+        qk_int8_sv_f16_accum_f16_attn_inst_buf as sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf,
+    )
     SM80_ENABLED = True
-except:
+except Exception:
     SM80_ENABLED = False
 
 try:
-    from . import sm89_compile
+    from .sm89_compile import (
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn,
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn,
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf,
+        qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf,
+    )
     SM89_ENABLED = True
-except:
+except Exception:
     SM89_ENABLED = False
 
 try:
-    from . import sm90_compile
+    from .sm90_compile import (
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90 as sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90,
+    )
     SM90_ENABLED = True
-except:
+except Exception:
     SM90_ENABLED = False
 
 from typing import Any, List, Literal, Optional, Tuple, Union
@@ -378,7 +390,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
 
     if pv_accum_dtype == "fp32":
         v = v.to(torch.float16)
-        lse = sm80_compile.qk_int8_sv_f16_accum_f32_attn(
+        lse = sm80_qk_int8_sv_f16_accum_f32_attn(
             q_int8,
             k_int8,
             v,
@@ -394,7 +406,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
     elif pv_accum_dtype == "fp16":
         if smooth_v:
             smoothed_v, vm = sub_mean(v, tensor_layout=tensor_layout)
-            lse = sm80_compile.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+            lse = sm80_qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 smoothed_v,
@@ -410,7 +422,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
         else:
             v = v.to(torch.float16)
-            lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn(
+            lse = sm80_qk_int8_sv_f16_accum_f16_attn(
                 q_int8,
                 k_int8,
                 v,
@@ -425,7 +437,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
     elif pv_accum_dtype == "fp16+fp32":
         v = v.to(torch.float16)
-        lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        lse = sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf(
             q_int8,
             k_int8,
             v,
@@ -650,7 +662,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
     )
     if pv_accum_dtype == "fp32":
         if smooth_v:
-            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+            lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -667,7 +679,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
         else:
-            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+            lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -683,7 +695,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp32":
-        lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -699,7 +711,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp16":
-        lse = sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        lse = sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -936,7 +948,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
 
     if pv_accum_dtype == "fp32":
         raise NotImplementedError("Please use pv_accum_dtype='fp32+fp32' for sm90.")
-        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        lse = sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
             q_int8,
             k_int8,
             v_fp8,
@@ -951,7 +963,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
             _return_lse,
         )
     elif pv_accum_dtype == "fp32+fp32":
-        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        lse = sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
             q_int8,
             k_int8,
             v_fp8,

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

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

build/torch29-cxx11-cu128-x86_64-linux/_sage_attention_6e51d70.abi3.so

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

build/torch29-cxx11-cu128-x86_64-linux/_sage_attention_cuda_4eabbf5.abi3.so

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

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

@@ -26,21 +26,33 @@ from .quant import per_channel_fp8
 from .quant_per_thread import per_thread_int8 as per_thread_int8_triton
 
 try:
-    from . import sm80_compile
+    from .sm80_compile import (
+        qk_int8_sv_f16_accum_f32_attn as sm80_qk_int8_sv_f16_accum_f32_attn,
+        qk_int8_sv_f16_accum_f16_fuse_v_mean_attn as sm80_qk_int8_sv_f16_accum_f16_fuse_v_mean_attn,
+        qk_int8_sv_f16_accum_f16_attn as sm80_qk_int8_sv_f16_accum_f16_attn,
+        qk_int8_sv_f16_accum_f16_attn_inst_buf as sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf,
+    )
     SM80_ENABLED = True
-except:
+except Exception:
     SM80_ENABLED = False
 
 try:
-    from . import sm89_compile
+    from .sm89_compile import (
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn,
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn,
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf,
+        qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf,
+    )
     SM89_ENABLED = True
-except:
+except Exception:
     SM89_ENABLED = False
 
 try:
-    from . import sm90_compile
+    from .sm90_compile import (
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90 as sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90,
+    )
     SM90_ENABLED = True
-except:
+except Exception:
     SM90_ENABLED = False
 
 from typing import Any, List, Literal, Optional, Tuple, Union
@@ -378,7 +390,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
 
     if pv_accum_dtype == "fp32":
         v = v.to(torch.float16)
-        lse = sm80_compile.qk_int8_sv_f16_accum_f32_attn(
+        lse = sm80_qk_int8_sv_f16_accum_f32_attn(
             q_int8,
             k_int8,
             v,
@@ -394,7 +406,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
     elif pv_accum_dtype == "fp16":
         if smooth_v:
             smoothed_v, vm = sub_mean(v, tensor_layout=tensor_layout)
-            lse = sm80_compile.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+            lse = sm80_qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 smoothed_v,
@@ -410,7 +422,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
         else:
             v = v.to(torch.float16)
-            lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn(
+            lse = sm80_qk_int8_sv_f16_accum_f16_attn(
                 q_int8,
                 k_int8,
                 v,
@@ -425,7 +437,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
     elif pv_accum_dtype == "fp16+fp32":
         v = v.to(torch.float16)
-        lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        lse = sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf(
             q_int8,
             k_int8,
             v,
@@ -650,7 +662,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
     )
     if pv_accum_dtype == "fp32":
         if smooth_v:
-            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+            lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -667,7 +679,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
         else:
-            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+            lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -683,7 +695,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp32":
-        lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -699,7 +711,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp16":
-        lse = sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        lse = sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -936,7 +948,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
 
     if pv_accum_dtype == "fp32":
         raise NotImplementedError("Please use pv_accum_dtype='fp32+fp32' for sm90.")
-        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        lse = sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
             q_int8,
             k_int8,
             v_fp8,
@@ -951,7 +963,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
             _return_lse,
         )
     elif pv_accum_dtype == "fp32+fp32":
-        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        lse = sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
             q_int8,
             k_int8,
             v_fp8,

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

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

build/torch29-cxx11-cu130-x86_64-linux/_sage_attention_6e51d70.abi3.so

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

build/torch29-cxx11-cu130-x86_64-linux/_sage_attention_cuda_4eabbf5.abi3.so

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

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

@@ -26,21 +26,33 @@ from .quant import per_channel_fp8
 from .quant_per_thread import per_thread_int8 as per_thread_int8_triton
 
 try:
-    from . import sm80_compile
+    from .sm80_compile import (
+        qk_int8_sv_f16_accum_f32_attn as sm80_qk_int8_sv_f16_accum_f32_attn,
+        qk_int8_sv_f16_accum_f16_fuse_v_mean_attn as sm80_qk_int8_sv_f16_accum_f16_fuse_v_mean_attn,
+        qk_int8_sv_f16_accum_f16_attn as sm80_qk_int8_sv_f16_accum_f16_attn,
+        qk_int8_sv_f16_accum_f16_attn_inst_buf as sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf,
+    )
     SM80_ENABLED = True
-except:
+except Exception:
     SM80_ENABLED = False
 
 try:
-    from . import sm89_compile
+    from .sm89_compile import (
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn,
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn,
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf,
+        qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf,
+    )
     SM89_ENABLED = True
-except:
+except Exception:
     SM89_ENABLED = False
 
 try:
-    from . import sm90_compile
+    from .sm90_compile import (
+        qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90 as sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90,
+    )
     SM90_ENABLED = True
-except:
+except Exception:
     SM90_ENABLED = False
 
 from typing import Any, List, Literal, Optional, Tuple, Union
@@ -378,7 +390,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
 
     if pv_accum_dtype == "fp32":
         v = v.to(torch.float16)
-        lse = sm80_compile.qk_int8_sv_f16_accum_f32_attn(
+        lse = sm80_qk_int8_sv_f16_accum_f32_attn(
             q_int8,
             k_int8,
             v,
@@ -394,7 +406,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
     elif pv_accum_dtype == "fp16":
         if smooth_v:
             smoothed_v, vm = sub_mean(v, tensor_layout=tensor_layout)
-            lse = sm80_compile.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+            lse = sm80_qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 smoothed_v,
@@ -410,7 +422,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
         else:
             v = v.to(torch.float16)
-            lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn(
+            lse = sm80_qk_int8_sv_f16_accum_f16_attn(
                 q_int8,
                 k_int8,
                 v,
@@ -425,7 +437,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
     elif pv_accum_dtype == "fp16+fp32":
         v = v.to(torch.float16)
-        lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        lse = sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf(
             q_int8,
             k_int8,
             v,
@@ -650,7 +662,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
     )
     if pv_accum_dtype == "fp32":
         if smooth_v:
-            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+            lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -667,7 +679,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
         else:
-            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+            lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -683,7 +695,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp32":
-        lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        lse = sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -699,7 +711,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp16":
-        lse = sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        lse = sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -936,7 +948,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
 
     if pv_accum_dtype == "fp32":
         raise NotImplementedError("Please use pv_accum_dtype='fp32+fp32' for sm90.")
-        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        lse = sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
             q_int8,
             k_int8,
             v_fp8,
@@ -951,7 +963,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
             _return_lse,
         )
     elif pv_accum_dtype == "fp32+fp32":
-        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        lse = sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
             q_int8,
             k_int8,
             v_fp8,