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,