Optimize triton_mrope with torch compile (#12112)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

python/sglang/srt/layers/rotary_embedding.py

@@ -1424,6 +1424,7 @@ class MRotaryEmbedding(RotaryEmbedding):
         else:
             return self._forward_native(positions, query, key)
 
+    @torch.compile(dynamic=True, backend=get_compiler_backend())
     def _forward_triton(
         self,
         positions: torch.Tensor,
@@ -1442,6 +1443,7 @@ class MRotaryEmbedding(RotaryEmbedding):
         if positions.ndim == 2:
             assert self.mrope_section
 
+            torch._dynamo.graph_break()
             q, k = triton_mrope(
                 query,
                 key,
@@ -1453,6 +1455,7 @@ class MRotaryEmbedding(RotaryEmbedding):
                 self.mrope_interleaved,
                 self.is_neox_style,
             )
+            torch._dynamo.graph_break()
 
             return q.reshape(query_shape), k.reshape(key_shape)