torch compile mrope

timmy-feng · timmy-feng · commit edafab3762af · 2025-08-20T02:06:31.000-04:00
diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py
@@ -1029,6 +1029,7 @@ def __init__(
                     f"Corrected mrope_section: {self.mrope_section} (sum={sum(self.mrope_section)})"
                 )
 
+    @torch.compile(dynamic=True)
     def forward(
         self,
         positions: torch.Tensor,

Original file line number	Diff line number	Diff line change
`@@ -1029,6 +1029,7 @@ def __init__(`
`1029`	`1029`	`f"Corrected mrope_section: {self.mrope_section} (sum={sum(self.mrope_section)})"`
`1030`	`1030`	`)`
`1031`	`1031`
	`1032`	`+ @torch.compile(dynamic=True)`
`1032`	`1033`	`def forward(`
`1033`	`1034`	`self,`
`1034`	`1035`	`positions: torch.Tensor,`