use default for torch.ops (#4835)

2025-03-27 19:09:58 -07:00
parent 10a9ab7b07
commit 31dfff7da7
7 changed files with 51 additions and 47 deletions
--- a/sgl-kernel/python/sgl_kernel/attention.py
+++ b/sgl-kernel/python/sgl_kernel/attention.py
@@ -2,6 +2,6 @@ import torch


 def lightning_attention_decode(q, k, v, past_kv, slope, output, new_kv):
-    torch.ops.sgl_kernel.lightning_attention_decode(
+    torch.ops.sgl_kernel.lightning_attention_decode.default(
        q, k, v, past_kv, slope, output, new_kv
    )