From 1158affa66b2f9d647069d1e231314e12fd40858 Mon Sep 17 00:00:00 2001
From: WateBear <540295877@qq.com>
Date: Wed, 26 Aug 2026 02:49:49 +0800
Subject: [PATCH] fix(mlu): make Sage attention compile safe

---
 .../ops/attn/cambricon_mlu/sage_attn.py       | 54 ++++++++++++++-----
 1 file changed, 42 insertions(+), 12 deletions(-)

diff --git a/lightx2v_platform/ops/attn/cambricon_mlu/sage_attn.py b/lightx2v_platform/ops/attn/cambricon_mlu/sage_attn.py
index bf540c32b..0cf790e62 100755
--- a/lightx2v_platform/ops/attn/cambricon_mlu/sage_attn.py
+++ b/lightx2v_platform/ops/attn/cambricon_mlu/sage_attn.py
@@ -11,6 +11,47 @@
     tmo = None
 
 
+@torch.library.custom_op(
+    "lightx2v::mlu_sage_attention",
+    mutates_args=(),
+    device_types="mlu",
+)
+def _mlu_sage_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_kv: int,
+    softmax_scale: float,
+    causal: bool,
+) -> torch.Tensor:
+    return tmo.sage_attn(
+        q=q,
+        k=k,
+        v=v,
+        cu_seq_lens_q=None,
+        cu_seq_lens_kv=None,
+        max_seq_len_q=max_seqlen_q,
+        max_seq_len_kv=max_seqlen_kv,
+        softmax_scale=softmax_scale,
+        is_causal=causal,
+        compute_dtype=torch.bfloat16,
+    )
+
+
+@_mlu_sage_attention.register_fake
+def _mlu_sage_attention_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_kv: int,
+    softmax_scale: float,
+    causal: bool,
+) -> torch.Tensor:
+    return v.new_empty((*q.shape[:-1], v.shape[-1]))
+
+
 @PLATFORM_ATTN_WEIGHT_REGISTER("mlu_sage_attn")
 class MluSageAttnWeight(AttnWeightTemplate):
     def __init__(self):
@@ -27,17 +68,6 @@ def apply(self, q, k, v, cu_seqlens_q=None, cu_seqlens_kv=None, max_seqlen_q=Non
         if softmax_scale is None:
             softmax_scale = 1 / math.sqrt(q.shape[-1])
         causal = kwds.get("causal", False)
-        x = tmo.sage_attn(
-            q=q,
-            k=k,
-            v=v,
-            cu_seq_lens_q=None,
-            cu_seq_lens_kv=None,
-            max_seq_len_kv=max_seqlen_kv,
-            max_seq_len_q=max_seqlen_q,
-            is_causal=causal,
-            compute_dtype=torch.bfloat16,
-            softmax_scale=softmax_scale,
-        )
+        x = _mlu_sage_attention(q, k, v, max_seqlen_q, max_seqlen_kv, softmax_scale, causal)
         x = x.reshape(bs * max_seqlen_q, -1)
         return x