Build uploaded using `kernels`.

.gitattributes

@@ -74,3 +74,4 @@ build/torch210-cxx11-cu130-x86_64-linux/_sage_attention_cuda_5568690.abi3.so fil
 build/torch29-cxx11-cu126-x86_64-linux/_sage_attention_cuda_5568690.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch29-cxx11-cu128-x86_64-linux/_sage_attention_cuda_5568690.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch29-cxx11-cu130-x86_64-linux/_sage_attention_cuda_5568690.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cu128-x86_64-windows/_sage_attention_cuda_a8f8348.pyd filter=lfs diff=lfs merge=lfs -text

build/torch210-cu128-x86_64-windows/_ops.py

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

build/torch210-cu128-x86_64-windows/{_sage_attention_cuda_554dbc8.pyd → _sage_attention_cuda_a8f8348.pyd}

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

build/torch210-cu128-x86_64-windows/core.py

@@ -16,6 +16,7 @@ limitations under the License.
 
 import torch
 import torch.nn.functional as F
+import warnings
 
 from ._ops import ops
 
@@ -33,8 +34,9 @@ try:
         qk_int8_sv_f16_accum_f16_attn_inst_buf as sm80_qk_int8_sv_f16_accum_f16_attn_inst_buf,
     )
     SM80_ENABLED = True
-except Exception:
+except Exception as e:
     SM80_ENABLED = False
+    warnings.warn(f"Failed to load SM80 SageAttention kernels: {e}")
 
 try:
     from .sm89_compile import (
@@ -44,20 +46,20 @@ try:
         qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf as sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf,
     )
     SM89_ENABLED = True
-except Exception:
+except Exception as e:
     SM89_ENABLED = False
+    warnings.warn(f"Failed to load SM89 SageAttention kernels: {e}")
 
 try:
     from .sm90_compile import (
         qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90 as sm90_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90,
     )
     SM90_ENABLED = True
-except Exception:
+except Exception as e:
     SM90_ENABLED = False
+    warnings.warn(f"Failed to load SM90 SageAttention kernels: {e}")
 
 from typing import Any, List, Literal, Optional, Tuple, Union
-import warnings
-
 
 import subprocess
 import re
@@ -148,6 +150,11 @@ def sageattn(
     """
     arch = get_cuda_arch_versions()[q.device.index]
     if arch == "sm80":
+        if not SM80_ENABLED:
+            raise RuntimeError(
+                "SM80 SageAttention kernels failed to load. "
+                "Ensure the kernel was compiled for SM80 (Ampere)."
+            )
         return sageattn_qk_int8_pv_fp16_cuda(
             q,
             k,
@@ -159,6 +166,11 @@ def sageattn(
             pv_accum_dtype="fp32",
         )
     elif arch == "sm89":
+        if not SM89_ENABLED:
+            raise RuntimeError(
+                "SM89 SageAttention kernels failed to load. "
+                "Ensure the kernel was compiled for SM89 (Ada Lovelace)."
+            )
         return sageattn_qk_int8_pv_fp8_cuda(
             q,
             k,
@@ -170,6 +182,11 @@ def sageattn(
             pv_accum_dtype="fp32+fp16",
         )
     elif arch == "sm90":
+        if not SM90_ENABLED:
+            raise RuntimeError(
+                "SM90 SageAttention kernels failed to load. "
+                "Ensure the kernel was compiled for SM90 (Hopper)."
+            )
         return sageattn_qk_int8_pv_fp8_cuda_sm90(
             q,
             k,
@@ -181,6 +198,11 @@ def sageattn(
             pv_accum_dtype="fp32+fp32",
         )
     elif arch == "sm120":
+        if not SM89_ENABLED:
+            raise RuntimeError(
+                "SM89 SageAttention kernels failed to load. "
+                "SM120 (Blackwell) uses SM89 kernels; ensure they were compiled."
+            )
         return sageattn_qk_int8_pv_fp8_cuda(
             q,
             k,

build/torch210-cu128-x86_64-windows/sm80_compile.py

@@ -2,148 +2,53 @@ 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:
+def _lse_fake_impl(query, tensor_layout, return_lse):
     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
+        return torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    return torch.empty((0))
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn"))
+def qk_int8_sv_f16_accum_f16_attn_fake(
+    query, key, value, output, query_scale, key_scale,
+    tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse,
+):
+    return _lse_fake_impl(query, tensor_layout, 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_f32_attn"))
+def qk_int8_sv_f16_accum_f32_attn_fake(
+    query, key, value, output, query_scale, key_scale,
+    tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse,
+):
+    return _lse_fake_impl(query, tensor_layout, return_lse)
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f16_accum_f16_attn_inst_buf"))
+def qk_int8_sv_f16_accum_f16_attn_inst_buf_fake(
+    query, key, value, output, query_scale, key_scale,
+    tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse,
+):
+    return _lse_fake_impl(query, tensor_layout, return_lse)
 
 
 @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
-    )
+def qk_int8_sv_f16_accum_f16_fuse_v_mean_attn_fake(
+    query, key, value, output, query_scale, key_scale, value_mean,
+    tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse,
+):
+    return _lse_fake_impl(query, tensor_layout, return_lse)
+
+
+qk_int8_sv_f16_accum_f16_attn = ops.qk_int8_sv_f16_accum_f16_attn
+qk_int8_sv_f16_accum_f32_attn = ops.qk_int8_sv_f16_accum_f32_attn
+qk_int8_sv_f16_accum_f16_attn_inst_buf = ops.qk_int8_sv_f16_accum_f16_attn_inst_buf
+qk_int8_sv_f16_accum_f16_fuse_v_mean_attn = ops.qk_int8_sv_f16_accum_f16_fuse_v_mean_attn

build/torch210-cu128-x86_64-windows/sm89_compile.py

@@ -2,145 +2,53 @@ 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:
+def _lse_fake_impl(query, tensor_layout, return_lse):
     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
+        return torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    return torch.empty((0))
+
 
+@torch.library.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn"))
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_fake(
+    query, key, value, output, query_scale, key_scale, value_scale,
+    tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse,
+):
+    return _lse_fake_impl(query, tensor_layout, 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.register_fake(add_op_namespace_prefix("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf"))
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_fake(
+    query, key, value, output, query_scale, key_scale, value_scale,
+    tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse,
+):
+    return _lse_fake_impl(query, tensor_layout, return_lse)
 
-@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_f16_fuse_v_scale_attn_inst_buf"))
+def qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf_fake(
+    query, key, value, output, query_scale, key_scale, value_scale,
+    tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse,
+):
+    return _lse_fake_impl(query, tensor_layout, 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
-    )
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn_fake(
+    query, key, value, output, query_scale, key_scale, value_scale, value_mean,
+    tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse,
+):
+    return _lse_fake_impl(query, tensor_layout, return_lse)
+
+
+qk_int8_sv_f8_accum_f32_fuse_v_scale_attn = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn
+qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf
+qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf = ops.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf
+qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn

build/torch210-cu128-x86_64-windows/sm90_compile.py

@@ -2,93 +2,35 @@ 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:
+def _lse_fake_impl(query, tensor_layout, return_lse):
     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
+        return torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device=query.device)
+    return torch.empty((0))
 
 
-@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_attn_inst_buf"))
+def qk_int8_sv_f8_accum_f32_attn_inst_buf_fake(
+    query, key, value, output, query_scale, key_scale,
+    tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse,
+):
+    return _lse_fake_impl(query, tensor_layout, 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
-    )
+def qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90_fake(
+    query, key, value, output, query_scale, key_scale, value_scale,
+    tensor_layout, is_causal, qk_quant_gran, sm_scale, return_lse,
+):
+    return _lse_fake_impl(query, tensor_layout, return_lse)
+
+
+qk_int8_sv_f8_accum_f32_attn_inst_buf = ops.qk_int8_sv_f8_accum_f32_attn_inst_buf
+qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90 = ops.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_sm90