From bb517fe39381b47b5fd304e105df29fb33754804 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Tue, 4 Nov 2025 23:26:56 +0800 Subject: [PATCH] [HotFix] Disable torch dynamo for mrope_triton kernel (#12593) Co-authored-by: luoyuan.luo --- python/sglang/srt/layers/rotary_embedding.py | 29 ++++++++++++++++++-- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index 64ac00ff1..5d0469991 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -1302,6 +1302,31 @@ def triton_mrope( return q, k +@torch._dynamo.disable() +def triton_mrope_wrapper( + query, + key, + cos, + sin, + mrope_section, + head_size, + rotary_dim, + mrope_interleaved, + is_neox_style, +): + return triton_mrope( + query, + key, + cos, + sin, + mrope_section, + head_size, + rotary_dim, + mrope_interleaved, + is_neox_style, + ) + + class MRotaryEmbedding(RotaryEmbedding): """Rotary Embedding with Multimodal Sections.""" @@ -1460,8 +1485,7 @@ class MRotaryEmbedding(RotaryEmbedding): if positions.ndim == 2: assert self.mrope_section - torch._dynamo.graph_break() - q, k = triton_mrope( + q, k = triton_mrope_wrapper( query, key, cos, @@ -1472,7 +1496,6 @@ class MRotaryEmbedding(RotaryEmbedding): self.mrope_interleaved, self.is_neox_style, ) - torch._dynamo.graph_break() return q.reshape(query_shape), k.reshape(key_shape)