minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Optimize triton_mrope with torch compile (#12112)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2025-10-27 23:49:22 +08:00
committed by GitHub
co-authored by luoyuan.luo
parent caa4819bfc
commit f389f01714
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/rotary_embedding.py
+3
View File
@@ -1424,6 +1424,7 @@ class MRotaryEmbedding(RotaryEmbedding):
else: else:
return self._forward_native(positions, query, key) return self._forward_native(positions, query, key)
@torch.compile(dynamic=True, backend=get_compiler_backend())
def _forward_triton( def _forward_triton(
self, self,
positions: torch.Tensor, positions: torch.Tensor,
@@ -1442,6 +1443,7 @@ class MRotaryEmbedding(RotaryEmbedding):
if positions.ndim == 2: if positions.ndim == 2:
assert self.mrope_section assert self.mrope_section
torch._dynamo.graph_break()
q, k = triton_mrope( q, k = triton_mrope(
query, query,
key, key,
@@ -1453,6 +1455,7 @@ class MRotaryEmbedding(RotaryEmbedding):
self.mrope_interleaved, self.mrope_interleaved,
self.is_neox_style, self.is_neox_style,
) )
torch._dynamo.graph_break()
return q.reshape(query_shape), k.reshape(key_shape) return q.reshape(query_shape), k.reshape(key_shape)