diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index 4b3856fd14b6..48f564f83436 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -1424,6 +1424,7 @@ def forward( 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 @@ def _forward_triton( if positions.ndim == 2: assert self.mrope_section + torch._dynamo.graph_break() q, k = triton_mrope( query, key, @@ -1453,6 +1455,7 @@ def _forward_triton( self.mrope_interleaved, self.is_neox_style, ) + torch._dynamo.graph_break() return q.reshape(query_shape), k.reshape(key_shape)