Build uploaded using `kernels`.

.gitattributes

@@ -42,3 +42,9 @@ build/torch28-cxx11-cu129-x86_64-linux/_sage_attention_fd27ff6.abi3.so filter=lf
 build/torch29-cxx11-cu126-x86_64-linux/_sage_attention_fd27ff6.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch29-cxx11-cu128-x86_64-linux/_sage_attention_fd27ff6.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch29-cxx11-cu130-x86_64-linux/_sage_attention_fd27ff6.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu126-x86_64-linux/_sage_attention_6e51d70.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu128-x86_64-linux/_sage_attention_6e51d70.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu130-x86_64-linux/_sage_attention_6e51d70.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-cxx11-cu126-x86_64-linux/_sage_attention_6e51d70.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-cxx11-cu128-x86_64-linux/_sage_attention_6e51d70.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-cxx11-cu130-x86_64-linux/_sage_attention_6e51d70.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_fd27ff6
-ops = torch.ops._sage_attention_fd27ff6
+from . import _sage_attention_6e51d70
+ops = torch.ops._sage_attention_6e51d70
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_sage_attention_fd27ff6::{op_name}"
+    return f"_sage_attention_6e51d70::{op_name}"

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

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

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

@@ -25,6 +25,24 @@ from .quant import sub_mean
 from .quant import per_channel_fp8
 from .quant_per_thread import per_thread_int8 as per_thread_int8_triton
 
+try:
+    from . import sm80_compile
+    SM80_ENABLED = True
+except:
+    SM80_ENABLED = False
+
+try:
+    from . import sm89_compile
+    SM89_ENABLED = True
+except:
+    SM89_ENABLED = False
+
+try:
+    from . import sm90_compile
+    SM90_ENABLED = True
+except:
+    SM90_ENABLED = False
+
 from typing import Any, List, Literal, Optional, Tuple, Union
 import warnings
 
@@ -116,7 +134,6 @@ def sageattn(
     - The tensors `q`, `k`, and `v` must have the dtype ``torch.float16`` or ``torch.bfloat16``
     - All tensors must be on the same cuda device.
     """
-    print("########tensor_layout", tensor_layout)
     arch = get_cuda_arch_versions()[q.device.index]
     if arch == "sm80":
         return sageattn_qk_int8_pv_fp16_cuda(
@@ -166,8 +183,6 @@ def sageattn(
     else:
         raise ValueError(f"Unsupported CUDA architecture: {arch}")
 
-
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp16_cuda(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -363,7 +378,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
 
     if pv_accum_dtype == "fp32":
         v = v.to(torch.float16)
-        lse = ops.qk_int8_sv_f16_accum_f32_attn(
+        lse = sm80_compile.qk_int8_sv_f16_accum_f32_attn(
             q_int8,
             k_int8,
             v,
@@ -379,7 +394,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 = ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+            lse = sm80_compile.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 smoothed_v,
@@ -395,7 +410,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
         else:
             v = v.to(torch.float16)
-            lse = ops.qk_int8_sv_f16_accum_f16_attn(
+            lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn(
                 q_int8,
                 k_int8,
                 v,
@@ -410,7 +425,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
     elif pv_accum_dtype == "fp16+fp32":
         v = v.to(torch.float16)
-        lse = ops.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn_inst_buf(
             q_int8,
             k_int8,
             v,
@@ -438,8 +453,6 @@ def sageattn_qk_int8_pv_fp16_cuda(
     else:
         return o
 
-
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp8_cuda(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -635,10 +648,9 @@ def sageattn_qk_int8_pv_fp8_cuda(
     v_fp8, v_scale, vm = per_channel_fp8(
         v, tensor_layout=tensor_layout, scale_max=quant_v_scale_max, smooth_v=smooth_v
     )
-    print("before kernel call")
     if pv_accum_dtype == "fp32":
         if smooth_v:
-            lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -655,7 +667,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
         else:
-            lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -671,7 +683,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp32":
-        lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -687,7 +699,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp16":
-        lse = ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        lse = sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -703,7 +715,6 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     o = o[..., :head_dim_og]
-    print("after kernel call")
     if return_lse:
         return (
             o,
@@ -715,7 +726,6 @@ def sageattn_qk_int8_pv_fp8_cuda(
         return o
 
 
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp8_cuda_sm90(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -926,7 +936,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 = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
             q_int8,
             k_int8,
             v_fp8,
@@ -941,7 +951,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
             _return_lse,
         )
     elif pv_accum_dtype == "fp32+fp32":
-        lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
             q_int8,
             k_int8,
             v_fp8,

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

@@ -0,0 +1,149 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_attn(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f32_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f32_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP32 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f32_attn(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_fuse_v_mean_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_mean: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+        query, key, value, output, query_scale, key_scale, value_mean,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+def sm80_qk_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"))(sm80_qk_fake_impl)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f32_attn"))(sm80_qk_fake_impl)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"))(sm80_qk_fake_impl)
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_fuse_v_mean_attn"))
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return sm80_qk_fake_impl(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

@@ -0,0 +1,146 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+def sm89_qk_with_key_value(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"))(sm89_qk_with_key_value)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf"))(sm89_qk_with_key_value)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"))(sm89_qk_with_key_value)
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        value_mean, tensor_layout, is_causal, qk_quant_gran, sm_scale,
+        return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn"))
+def sm89_qk_with_key_value_mean(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return sm89_qk_with_key_value(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

@@ -0,0 +1,94 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_attn_inst_buf"))
+def qk_int8_sv_f8_accum_f32_attn_inst_buf_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90"))
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return qk_int8_sv_f8_accum_f32_attn_inst_buf_fake_impl(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

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

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

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

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

@@ -25,6 +25,24 @@ from .quant import sub_mean
 from .quant import per_channel_fp8
 from .quant_per_thread import per_thread_int8 as per_thread_int8_triton
 
+try:
+    from . import sm80_compile
+    SM80_ENABLED = True
+except:
+    SM80_ENABLED = False
+
+try:
+    from . import sm89_compile
+    SM89_ENABLED = True
+except:
+    SM89_ENABLED = False
+
+try:
+    from . import sm90_compile
+    SM90_ENABLED = True
+except:
+    SM90_ENABLED = False
+
 from typing import Any, List, Literal, Optional, Tuple, Union
 import warnings
 
@@ -116,7 +134,6 @@ def sageattn(
     - The tensors `q`, `k`, and `v` must have the dtype ``torch.float16`` or ``torch.bfloat16``
     - All tensors must be on the same cuda device.
     """
-    print("########tensor_layout", tensor_layout)
     arch = get_cuda_arch_versions()[q.device.index]
     if arch == "sm80":
         return sageattn_qk_int8_pv_fp16_cuda(
@@ -166,8 +183,6 @@ def sageattn(
     else:
         raise ValueError(f"Unsupported CUDA architecture: {arch}")
 
-
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp16_cuda(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -363,7 +378,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
 
     if pv_accum_dtype == "fp32":
         v = v.to(torch.float16)
-        lse = ops.qk_int8_sv_f16_accum_f32_attn(
+        lse = sm80_compile.qk_int8_sv_f16_accum_f32_attn(
             q_int8,
             k_int8,
             v,
@@ -379,7 +394,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 = ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+            lse = sm80_compile.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 smoothed_v,
@@ -395,7 +410,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
         else:
             v = v.to(torch.float16)
-            lse = ops.qk_int8_sv_f16_accum_f16_attn(
+            lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn(
                 q_int8,
                 k_int8,
                 v,
@@ -410,7 +425,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
     elif pv_accum_dtype == "fp16+fp32":
         v = v.to(torch.float16)
-        lse = ops.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn_inst_buf(
             q_int8,
             k_int8,
             v,
@@ -438,8 +453,6 @@ def sageattn_qk_int8_pv_fp16_cuda(
     else:
         return o
 
-
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp8_cuda(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -635,10 +648,9 @@ def sageattn_qk_int8_pv_fp8_cuda(
     v_fp8, v_scale, vm = per_channel_fp8(
         v, tensor_layout=tensor_layout, scale_max=quant_v_scale_max, smooth_v=smooth_v
     )
-    print("before kernel call")
     if pv_accum_dtype == "fp32":
         if smooth_v:
-            lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -655,7 +667,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
         else:
-            lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -671,7 +683,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp32":
-        lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -687,7 +699,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp16":
-        lse = ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        lse = sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -703,7 +715,6 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     o = o[..., :head_dim_og]
-    print("after kernel call")
     if return_lse:
         return (
             o,
@@ -715,7 +726,6 @@ def sageattn_qk_int8_pv_fp8_cuda(
         return o
 
 
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp8_cuda_sm90(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -926,7 +936,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 = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
             q_int8,
             k_int8,
             v_fp8,
@@ -941,7 +951,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
             _return_lse,
         )
     elif pv_accum_dtype == "fp32+fp32":
-        lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        lse = sm90_compile.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/sm80_compile.py

@@ -0,0 +1,149 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_attn(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f32_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f32_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP32 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f32_attn(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_fuse_v_mean_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_mean: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+        query, key, value, output, query_scale, key_scale, value_mean,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+def sm80_qk_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"))(sm80_qk_fake_impl)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f32_attn"))(sm80_qk_fake_impl)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"))(sm80_qk_fake_impl)
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_fuse_v_mean_attn"))
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return sm80_qk_fake_impl(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

@@ -0,0 +1,146 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+def sm89_qk_with_key_value(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"))(sm89_qk_with_key_value)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf"))(sm89_qk_with_key_value)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"))(sm89_qk_with_key_value)
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        value_mean, tensor_layout, is_causal, qk_quant_gran, sm_scale,
+        return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn"))
+def sm89_qk_with_key_value_mean(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return sm89_qk_with_key_value(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

@@ -0,0 +1,94 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_attn_inst_buf"))
+def qk_int8_sv_f8_accum_f32_attn_inst_buf_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90"))
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return qk_int8_sv_f8_accum_f32_attn_inst_buf_fake_impl(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

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

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

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

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

@@ -25,6 +25,24 @@ from .quant import sub_mean
 from .quant import per_channel_fp8
 from .quant_per_thread import per_thread_int8 as per_thread_int8_triton
 
+try:
+    from . import sm80_compile
+    SM80_ENABLED = True
+except:
+    SM80_ENABLED = False
+
+try:
+    from . import sm89_compile
+    SM89_ENABLED = True
+except:
+    SM89_ENABLED = False
+
+try:
+    from . import sm90_compile
+    SM90_ENABLED = True
+except:
+    SM90_ENABLED = False
+
 from typing import Any, List, Literal, Optional, Tuple, Union
 import warnings
 
@@ -116,7 +134,6 @@ def sageattn(
     - The tensors `q`, `k`, and `v` must have the dtype ``torch.float16`` or ``torch.bfloat16``
     - All tensors must be on the same cuda device.
     """
-    print("########tensor_layout", tensor_layout)
     arch = get_cuda_arch_versions()[q.device.index]
     if arch == "sm80":
         return sageattn_qk_int8_pv_fp16_cuda(
@@ -166,8 +183,6 @@ def sageattn(
     else:
         raise ValueError(f"Unsupported CUDA architecture: {arch}")
 
-
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp16_cuda(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -363,7 +378,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
 
     if pv_accum_dtype == "fp32":
         v = v.to(torch.float16)
-        lse = ops.qk_int8_sv_f16_accum_f32_attn(
+        lse = sm80_compile.qk_int8_sv_f16_accum_f32_attn(
             q_int8,
             k_int8,
             v,
@@ -379,7 +394,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 = ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+            lse = sm80_compile.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 smoothed_v,
@@ -395,7 +410,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
         else:
             v = v.to(torch.float16)
-            lse = ops.qk_int8_sv_f16_accum_f16_attn(
+            lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn(
                 q_int8,
                 k_int8,
                 v,
@@ -410,7 +425,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
     elif pv_accum_dtype == "fp16+fp32":
         v = v.to(torch.float16)
-        lse = ops.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn_inst_buf(
             q_int8,
             k_int8,
             v,
@@ -438,8 +453,6 @@ def sageattn_qk_int8_pv_fp16_cuda(
     else:
         return o
 
-
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp8_cuda(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -635,10 +648,9 @@ def sageattn_qk_int8_pv_fp8_cuda(
     v_fp8, v_scale, vm = per_channel_fp8(
         v, tensor_layout=tensor_layout, scale_max=quant_v_scale_max, smooth_v=smooth_v
     )
-    print("before kernel call")
     if pv_accum_dtype == "fp32":
         if smooth_v:
-            lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -655,7 +667,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
         else:
-            lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -671,7 +683,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp32":
-        lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -687,7 +699,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp16":
-        lse = ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        lse = sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -703,7 +715,6 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     o = o[..., :head_dim_og]
-    print("after kernel call")
     if return_lse:
         return (
             o,
@@ -715,7 +726,6 @@ def sageattn_qk_int8_pv_fp8_cuda(
         return o
 
 
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp8_cuda_sm90(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -926,7 +936,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 = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
             q_int8,
             k_int8,
             v_fp8,
@@ -941,7 +951,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
             _return_lse,
         )
     elif pv_accum_dtype == "fp32+fp32":
-        lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        lse = sm90_compile.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/sm80_compile.py

@@ -0,0 +1,149 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_attn(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f32_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f32_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP32 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f32_attn(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_fuse_v_mean_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_mean: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+        query, key, value, output, query_scale, key_scale, value_mean,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+def sm80_qk_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"))(sm80_qk_fake_impl)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f32_attn"))(sm80_qk_fake_impl)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"))(sm80_qk_fake_impl)
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_fuse_v_mean_attn"))
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return sm80_qk_fake_impl(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

@@ -0,0 +1,146 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+def sm89_qk_with_key_value(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"))(sm89_qk_with_key_value)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf"))(sm89_qk_with_key_value)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"))(sm89_qk_with_key_value)
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        value_mean, tensor_layout, is_causal, qk_quant_gran, sm_scale,
+        return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn"))
+def sm89_qk_with_key_value_mean(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return sm89_qk_with_key_value(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

@@ -0,0 +1,94 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_attn_inst_buf"))
+def qk_int8_sv_f8_accum_f32_attn_inst_buf_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90"))
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return qk_int8_sv_f8_accum_f32_attn_inst_buf_fake_impl(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

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

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

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

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

@@ -25,6 +25,24 @@ from .quant import sub_mean
 from .quant import per_channel_fp8
 from .quant_per_thread import per_thread_int8 as per_thread_int8_triton
 
+try:
+    from . import sm80_compile
+    SM80_ENABLED = True
+except:
+    SM80_ENABLED = False
+
+try:
+    from . import sm89_compile
+    SM89_ENABLED = True
+except:
+    SM89_ENABLED = False
+
+try:
+    from . import sm90_compile
+    SM90_ENABLED = True
+except:
+    SM90_ENABLED = False
+
 from typing import Any, List, Literal, Optional, Tuple, Union
 import warnings
 
@@ -116,7 +134,6 @@ def sageattn(
     - The tensors `q`, `k`, and `v` must have the dtype ``torch.float16`` or ``torch.bfloat16``
     - All tensors must be on the same cuda device.
     """
-    print("########tensor_layout", tensor_layout)
     arch = get_cuda_arch_versions()[q.device.index]
     if arch == "sm80":
         return sageattn_qk_int8_pv_fp16_cuda(
@@ -166,8 +183,6 @@ def sageattn(
     else:
         raise ValueError(f"Unsupported CUDA architecture: {arch}")
 
-
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp16_cuda(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -363,7 +378,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
 
     if pv_accum_dtype == "fp32":
         v = v.to(torch.float16)
-        lse = ops.qk_int8_sv_f16_accum_f32_attn(
+        lse = sm80_compile.qk_int8_sv_f16_accum_f32_attn(
             q_int8,
             k_int8,
             v,
@@ -379,7 +394,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 = ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+            lse = sm80_compile.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 smoothed_v,
@@ -395,7 +410,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
         else:
             v = v.to(torch.float16)
-            lse = ops.qk_int8_sv_f16_accum_f16_attn(
+            lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn(
                 q_int8,
                 k_int8,
                 v,
@@ -410,7 +425,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
     elif pv_accum_dtype == "fp16+fp32":
         v = v.to(torch.float16)
-        lse = ops.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn_inst_buf(
             q_int8,
             k_int8,
             v,
@@ -438,8 +453,6 @@ def sageattn_qk_int8_pv_fp16_cuda(
     else:
         return o
 
-
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp8_cuda(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -635,10 +648,9 @@ def sageattn_qk_int8_pv_fp8_cuda(
     v_fp8, v_scale, vm = per_channel_fp8(
         v, tensor_layout=tensor_layout, scale_max=quant_v_scale_max, smooth_v=smooth_v
     )
-    print("before kernel call")
     if pv_accum_dtype == "fp32":
         if smooth_v:
-            lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -655,7 +667,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
         else:
-            lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -671,7 +683,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp32":
-        lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -687,7 +699,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp16":
-        lse = ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        lse = sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -703,7 +715,6 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     o = o[..., :head_dim_og]
-    print("after kernel call")
     if return_lse:
         return (
             o,
@@ -715,7 +726,6 @@ def sageattn_qk_int8_pv_fp8_cuda(
         return o
 
 
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp8_cuda_sm90(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -926,7 +936,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 = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
             q_int8,
             k_int8,
             v_fp8,
@@ -941,7 +951,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
             _return_lse,
         )
     elif pv_accum_dtype == "fp32+fp32":
-        lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        lse = sm90_compile.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/sm80_compile.py

@@ -0,0 +1,149 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_attn(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f32_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f32_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP32 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f32_attn(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_fuse_v_mean_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_mean: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+        query, key, value, output, query_scale, key_scale, value_mean,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+def sm80_qk_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"))(sm80_qk_fake_impl)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f32_attn"))(sm80_qk_fake_impl)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"))(sm80_qk_fake_impl)
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_fuse_v_mean_attn"))
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return sm80_qk_fake_impl(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

@@ -0,0 +1,146 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+def sm89_qk_with_key_value(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"))(sm89_qk_with_key_value)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf"))(sm89_qk_with_key_value)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"))(sm89_qk_with_key_value)
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        value_mean, tensor_layout, is_causal, qk_quant_gran, sm_scale,
+        return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn"))
+def sm89_qk_with_key_value_mean(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return sm89_qk_with_key_value(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

@@ -0,0 +1,94 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_attn_inst_buf"))
+def qk_int8_sv_f8_accum_f32_attn_inst_buf_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90"))
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return qk_int8_sv_f8_accum_f32_attn_inst_buf_fake_impl(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

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

build/torch29-cxx11-cu128-x86_64-linux/{_sage_attention_fd27ff6.abi3.so → _sage_attention_6e51d70.abi3.so}

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

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

@@ -25,6 +25,24 @@ from .quant import sub_mean
 from .quant import per_channel_fp8
 from .quant_per_thread import per_thread_int8 as per_thread_int8_triton
 
+try:
+    from . import sm80_compile
+    SM80_ENABLED = True
+except:
+    SM80_ENABLED = False
+
+try:
+    from . import sm89_compile
+    SM89_ENABLED = True
+except:
+    SM89_ENABLED = False
+
+try:
+    from . import sm90_compile
+    SM90_ENABLED = True
+except:
+    SM90_ENABLED = False
+
 from typing import Any, List, Literal, Optional, Tuple, Union
 import warnings
 
@@ -116,7 +134,6 @@ def sageattn(
     - The tensors `q`, `k`, and `v` must have the dtype ``torch.float16`` or ``torch.bfloat16``
     - All tensors must be on the same cuda device.
     """
-    print("########tensor_layout", tensor_layout)
     arch = get_cuda_arch_versions()[q.device.index]
     if arch == "sm80":
         return sageattn_qk_int8_pv_fp16_cuda(
@@ -166,8 +183,6 @@ def sageattn(
     else:
         raise ValueError(f"Unsupported CUDA architecture: {arch}")
 
-
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp16_cuda(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -363,7 +378,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
 
     if pv_accum_dtype == "fp32":
         v = v.to(torch.float16)
-        lse = ops.qk_int8_sv_f16_accum_f32_attn(
+        lse = sm80_compile.qk_int8_sv_f16_accum_f32_attn(
             q_int8,
             k_int8,
             v,
@@ -379,7 +394,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 = ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+            lse = sm80_compile.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 smoothed_v,
@@ -395,7 +410,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
         else:
             v = v.to(torch.float16)
-            lse = ops.qk_int8_sv_f16_accum_f16_attn(
+            lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn(
                 q_int8,
                 k_int8,
                 v,
@@ -410,7 +425,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
     elif pv_accum_dtype == "fp16+fp32":
         v = v.to(torch.float16)
-        lse = ops.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn_inst_buf(
             q_int8,
             k_int8,
             v,
@@ -438,8 +453,6 @@ def sageattn_qk_int8_pv_fp16_cuda(
     else:
         return o
 
-
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp8_cuda(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -635,10 +648,9 @@ def sageattn_qk_int8_pv_fp8_cuda(
     v_fp8, v_scale, vm = per_channel_fp8(
         v, tensor_layout=tensor_layout, scale_max=quant_v_scale_max, smooth_v=smooth_v
     )
-    print("before kernel call")
     if pv_accum_dtype == "fp32":
         if smooth_v:
-            lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -655,7 +667,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
         else:
-            lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -671,7 +683,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp32":
-        lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -687,7 +699,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp16":
-        lse = ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        lse = sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -703,7 +715,6 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     o = o[..., :head_dim_og]
-    print("after kernel call")
     if return_lse:
         return (
             o,
@@ -715,7 +726,6 @@ def sageattn_qk_int8_pv_fp8_cuda(
         return o
 
 
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp8_cuda_sm90(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -926,7 +936,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 = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
             q_int8,
             k_int8,
             v_fp8,
@@ -941,7 +951,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
             _return_lse,
         )
     elif pv_accum_dtype == "fp32+fp32":
-        lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        lse = sm90_compile.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/sm80_compile.py

@@ -0,0 +1,149 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_attn(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f32_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f32_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP32 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f32_attn(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_fuse_v_mean_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_mean: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+        query, key, value, output, query_scale, key_scale, value_mean,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+def sm80_qk_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"))(sm80_qk_fake_impl)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f32_attn"))(sm80_qk_fake_impl)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"))(sm80_qk_fake_impl)
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_fuse_v_mean_attn"))
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return sm80_qk_fake_impl(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

@@ -0,0 +1,146 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+def sm89_qk_with_key_value(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"))(sm89_qk_with_key_value)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf"))(sm89_qk_with_key_value)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"))(sm89_qk_with_key_value)
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        value_mean, tensor_layout, is_causal, qk_quant_gran, sm_scale,
+        return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn"))
+def sm89_qk_with_key_value_mean(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return sm89_qk_with_key_value(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

@@ -0,0 +1,94 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_attn_inst_buf"))
+def qk_int8_sv_f8_accum_f32_attn_inst_buf_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90"))
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return qk_int8_sv_f8_accum_f32_attn_inst_buf_fake_impl(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

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

build/torch29-cxx11-cu130-x86_64-linux/{_sage_attention_fd27ff6.abi3.so → _sage_attention_6e51d70.abi3.so}

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

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

@@ -25,6 +25,24 @@ from .quant import sub_mean
 from .quant import per_channel_fp8
 from .quant_per_thread import per_thread_int8 as per_thread_int8_triton
 
+try:
+    from . import sm80_compile
+    SM80_ENABLED = True
+except:
+    SM80_ENABLED = False
+
+try:
+    from . import sm89_compile
+    SM89_ENABLED = True
+except:
+    SM89_ENABLED = False
+
+try:
+    from . import sm90_compile
+    SM90_ENABLED = True
+except:
+    SM90_ENABLED = False
+
 from typing import Any, List, Literal, Optional, Tuple, Union
 import warnings
 
@@ -116,7 +134,6 @@ def sageattn(
     - The tensors `q`, `k`, and `v` must have the dtype ``torch.float16`` or ``torch.bfloat16``
     - All tensors must be on the same cuda device.
     """
-    print("########tensor_layout", tensor_layout)
     arch = get_cuda_arch_versions()[q.device.index]
     if arch == "sm80":
         return sageattn_qk_int8_pv_fp16_cuda(
@@ -166,8 +183,6 @@ def sageattn(
     else:
         raise ValueError(f"Unsupported CUDA architecture: {arch}")
 
-
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp16_cuda(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -363,7 +378,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
 
     if pv_accum_dtype == "fp32":
         v = v.to(torch.float16)
-        lse = ops.qk_int8_sv_f16_accum_f32_attn(
+        lse = sm80_compile.qk_int8_sv_f16_accum_f32_attn(
             q_int8,
             k_int8,
             v,
@@ -379,7 +394,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 = ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+            lse = sm80_compile.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 smoothed_v,
@@ -395,7 +410,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
         else:
             v = v.to(torch.float16)
-            lse = ops.qk_int8_sv_f16_accum_f16_attn(
+            lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn(
                 q_int8,
                 k_int8,
                 v,
@@ -410,7 +425,7 @@ def sageattn_qk_int8_pv_fp16_cuda(
             )
     elif pv_accum_dtype == "fp16+fp32":
         v = v.to(torch.float16)
-        lse = ops.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        lse = sm80_compile.qk_int8_sv_f16_accum_f16_attn_inst_buf(
             q_int8,
             k_int8,
             v,
@@ -438,8 +453,6 @@ def sageattn_qk_int8_pv_fp16_cuda(
     else:
         return o
 
-
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp8_cuda(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -635,10 +648,9 @@ def sageattn_qk_int8_pv_fp8_cuda(
     v_fp8, v_scale, vm = per_channel_fp8(
         v, tensor_layout=tensor_layout, scale_max=quant_v_scale_max, smooth_v=smooth_v
     )
-    print("before kernel call")
     if pv_accum_dtype == "fp32":
         if smooth_v:
-            lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -655,7 +667,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
         else:
-            lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+            lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
                 q_int8,
                 k_int8,
                 v_fp8,
@@ -671,7 +683,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
             )
             torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp32":
-        lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        lse = sm89_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -687,7 +699,7 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     elif pv_accum_dtype == "fp32+fp16":
-        lse = ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        lse = sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
             q_int8,
             k_int8,
             v_fp8,
@@ -703,7 +715,6 @@ def sageattn_qk_int8_pv_fp8_cuda(
         )
         torch.cuda.synchronize()
     o = o[..., :head_dim_og]
-    print("after kernel call")
     if return_lse:
         return (
             o,
@@ -715,7 +726,6 @@ def sageattn_qk_int8_pv_fp8_cuda(
         return o
 
 
-@torch.compiler.disable
 def sageattn_qk_int8_pv_fp8_cuda_sm90(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -926,7 +936,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 = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        lse = sm90_compile.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
             q_int8,
             k_int8,
             v_fp8,
@@ -941,7 +951,7 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
             _return_lse,
         )
     elif pv_accum_dtype == "fp32+fp32":
-        lse = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        lse = sm90_compile.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/sm80_compile.py

@@ -0,0 +1,149 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_attn(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f32_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f32_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP32 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f32_attn(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_fuse_v_mean_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_mean: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    """
+    Custom CUDA kernel for SageAttention with INT8 quantization for Q and K, FP16 PV with FP16 accumulation.
+    """
+    return ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(
+        query, key, value, output, query_scale, key_scale, value_mean,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+def sm80_qk_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"))(sm80_qk_fake_impl)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f32_attn"))(sm80_qk_fake_impl)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"))(sm80_qk_fake_impl)
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_fuse_v_mean_attn"))
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return sm80_qk_fake_impl(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

@@ -0,0 +1,146 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+def sm89_qk_with_key_value(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"))(sm89_qk_with_key_value)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf"))(sm89_qk_with_key_value)
+torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"))(sm89_qk_with_key_value)
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        value_mean, tensor_layout, is_causal, qk_quant_gran, sm_scale,
+        return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn"))
+def sm89_qk_with_key_value_mean(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    value_mean: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return sm89_qk_with_key_value(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )

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

@@ -0,0 +1,94 @@
+from ._ops import ops
+import torch
+from ._ops import add_op_namespace_prefix
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_attn_inst_buf"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_attn_inst_buf(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_attn_inst_buf(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_attn_inst_buf"))
+def qk_int8_sv_f8_accum_f32_attn_inst_buf_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    batch_size = query.size(0)
+
+    if tensor_layout == 0:
+        num_qo_heads = query.size(2)
+        qo_len = query.size(1)
+    else:
+        num_qo_heads = query.size(1)
+        qo_len = query.size(2)
+
+    if return_lse:
+        lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    else:
+        lse = torch.empty((0))
+    return lse
+
+
+@torch.library.custom_op(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90"), mutates_args=("output",), device_types="cuda")
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor,
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90(
+        query, key, value, output, query_scale, key_scale, value_scale,
+        tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90"))
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90_fake_impl(
+    query: torch.Tensor, 
+    key: torch.Tensor, 
+    value: torch.Tensor, 
+    output: torch.Tensor, 
+    query_scale: torch.Tensor, 
+    key_scale: torch.Tensor, 
+    value_scale: torch.Tensor, 
+    tensor_layout: int, 
+    is_causal: int, 
+    qk_quant_gran: int, 
+    sm_scale: float,
+    return_lse: int,
+) -> torch.Tensor:
+    return qk_int8_sv_f8_accum_f32_attn_inst_buf_fake_impl(
+        query, key, value, output, query_scale, key_scale, tensor_layout,
+        is_causal, qk_quant_gran, sm_scale, return_lse
+    )