From 6489f77733816fcc1a70e75d09c14e12cb8a299b Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Date: Mon, 16 Mar 2026 23:14:27 +0800
Subject: [PATCH] [Diffusion] Fix compile graph broken by flashinfer rope
 (#20699)

---
 .../runtime/layers/rotary_embedding/utils.py  | 28 ++++++++++++++++---
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py
index 4eaa106e8..57af17e52 100644
--- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py
+++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py
@@ -5,6 +5,28 @@ from typing import Optional, Tuple
 import torch
 
 from sglang.jit_kernel.diffusion.triton.rotary import apply_rotary_embedding
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.srt.utils.custom_op import register_custom_op_from_extern
+
+_is_cuda = current_platform.is_cuda()
+if _is_cuda:
+    try:
+        from flashinfer.rope import (
+            apply_rope_with_cos_sin_cache_inplace as _flashinfer_apply_rope_inplace,
+        )
+    except Exception:
+        _flashinfer_apply_rope_inplace = None
+else:
+    _flashinfer_apply_rope_inplace = None
+
+if _flashinfer_apply_rope_inplace is not None:
+    flashinfer_apply_rope_inplace = register_custom_op_from_extern(
+        _flashinfer_apply_rope_inplace,
+        op_name="flashinfer_apply_rope_with_cos_sin_cache_inplace",
+        mutates_args=["query", "key"],
+    )
+else:
+    flashinfer_apply_rope_inplace = None
 
 
 def _apply_rotary_emb(
@@ -67,9 +89,7 @@ def apply_flashinfer_rope_qk_inplace(
     if head_size != d:
         raise ValueError(f"head_size mismatch: inferred {d}, but head_size={head_size}")
 
-    try:
-        from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
-    except ImportError:
+    if flashinfer_apply_rope_inplace is None:
         # Triton fallback for AMD/ROCm where FlashInfer is not available
         import warnings
 
@@ -110,7 +130,7 @@ def apply_flashinfer_rope_qk_inplace(
 
     q_flat = q.reshape(bsz * seqlen, nheads * d).contiguous()
     k_flat = k.reshape(bsz * seqlen, nheads * d).contiguous()
-    apply_rope_with_cos_sin_cache_inplace(
+    flashinfer_apply_rope_inplace(
         positions=positions,
         query=q_flat,
         key=k_flat,