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

[VLM] Support Piecewise CUDA Graph for Qwen2.5-VL (#13055)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Co-authored-by: Yuhao Yang <yhyang201@gmail.com>
This commit is contained in:
Yuan Luo
2025-11-20 10:23:44 +08:00
committed by GitHub
co-authored by luoyuan.luo Yuhao Yang
parent 67fca6b297
commit af6bcadcf7
10 changed files with 710 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/rotary_embedding.py
+8 -5
View File
@@ -1301,7 +1301,6 @@ def triton_mrope(
return q, k
@torch._dynamo.disable()
def triton_mrope_wrapper(
query,
key,
@@ -1428,15 +1427,18 @@ class MRotaryEmbedding(RotaryEmbedding):
dim=-1,
)
seq_len_q = query.shape[0]
query_shape = query.shape
query = query.view(num_tokens, -1, self.head_size)
query = query.view(seq_len_q, -1, self.head_size)
query_rot = query[..., : self.rotary_dim]
query_pass = query[..., self.rotary_dim :]
query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
seq_len_k = key.shape[0]
key_shape = key.shape
key = key.view(num_tokens, -1, self.head_size)
key = key.view(seq_len_k, -1, self.head_size)
key_rot = key[..., : self.rotary_dim]
key_pass = key[..., self.rotary_dim :]
key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
@@ -1467,7 +1469,6 @@ 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,
@@ -1502,7 +1503,9 @@ class MRotaryEmbedding(RotaryEmbedding):
return q.reshape(query_shape), k.reshape(key_shape)
query = query.view(num_tokens, -1, self.head_size)
seq_len_q = query.shape[0]
query = query.view(seq_len_q, -1, self.head_size)
query_rot = query[..., : self.rotary_dim]
query_pass = query[..., self.rotary_dim :]
query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)