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

[Diffusion] Avoid cpu2gpu sync in flashinfer rope and apply flashinfer rope to wanvideo (#16668)

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Xiaoyu Zhang
2026-01-08 22:44:38 +08:00
committed by GitHub
co-authored by Mick
parent f52ae586b6
commit 294ff71d18
2 changed files with 32 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
+1 -3
View File
@@ -94,7 +94,7 @@ def apply_flashinfer_rope_qk_inplace(
return q_rot.view(bsz, seqlen, nheads, d), k_rot.view(bsz, seqlen, nheads, d) return q_rot.view(bsz, seqlen, nheads, d), k_rot.view(bsz, seqlen, nheads, d)
if positions is None: if positions is None:
pos_1d = torch.arange(seqlen, device="cpu", dtype=torch.long) pos_1d = torch.arange(seqlen, device=q.device, dtype=torch.long)
positions = pos_1d if bsz == 1 else pos_1d.repeat(bsz) positions = pos_1d if bsz == 1 else pos_1d.repeat(bsz)
else: else:
if not ( if not (
@@ -108,8 +108,6 @@ def apply_flashinfer_rope_qk_inplace(
f"positions length must be bsz*seqlen={bsz*seqlen}, got {positions.numel()}" f"positions length must be bsz*seqlen={bsz*seqlen}, got {positions.numel()}"
) )
positions = positions.to(q.device, non_blocking=True)
q_flat = q.reshape(bsz * seqlen, nheads * d).contiguous() q_flat = q.reshape(bsz * seqlen, nheads * d).contiguous()
k_flat = k.reshape(bsz * seqlen, nheads * d).contiguous() k_flat = k.reshape(bsz * seqlen, nheads * d).contiguous()
apply_rope_with_cos_sin_cache_inplace( apply_rope_with_cos_sin_cache_inplace(
python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
+31 -6
View File
@@ -30,6 +30,7 @@ from sglang.multimodal_gen.runtime.layers.mlp import MLP
from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
NDRotaryEmbedding, NDRotaryEmbedding,
_apply_rotary_emb, _apply_rotary_emb,
apply_flashinfer_rope_qk_inplace,
) )
from sglang.multimodal_gen.runtime.layers.visual_embedding import ( from sglang.multimodal_gen.runtime.layers.visual_embedding import (
ModulateProjection, ModulateProjection,
@@ -399,9 +400,21 @@ class WanTransformerBlock(nn.Module):
# Apply rotary embeddings # Apply rotary embeddings
cos, sin = freqs_cis cos, sin = freqs_cis
query, key = _apply_rotary_emb( if query.is_cuda and query.shape == key.shape:
query, cos, sin, is_neox_style=False cos_sin_cache = torch.cat(
), _apply_rotary_emb(key, cos, sin, is_neox_style=False) [
cos.to(dtype=torch.float32).contiguous(),
sin.to(dtype=torch.float32).contiguous(),
],
dim=-1,
)
query, key = apply_flashinfer_rope_qk_inplace(
query, key, cos_sin_cache, is_neox=False
)
else:
query, key = _apply_rotary_emb(
query, cos, sin, is_neox_style=False
), _apply_rotary_emb(key, cos, sin, is_neox_style=False)
attn_output = self.attn1(query, key, value) attn_output = self.attn1(query, key, value)
attn_output = attn_output.flatten(2) attn_output = attn_output.flatten(2)
attn_output, _ = self.to_out(attn_output) attn_output, _ = self.to_out(attn_output)
@@ -569,9 +582,21 @@ class WanTransformerBlock_VSA(nn.Module):
# Apply rotary embeddings # Apply rotary embeddings
cos, sin = freqs_cis cos, sin = freqs_cis
query, key = _apply_rotary_emb( if query.is_cuda and query.shape == key.shape:
query, cos, sin, is_neox_style=False cos_sin_cache = torch.cat(
), _apply_rotary_emb(key, cos, sin, is_neox_style=False) [
cos.to(dtype=torch.float32).contiguous(),
sin.to(dtype=torch.float32).contiguous(),
],
dim=-1,
)
query, key = apply_flashinfer_rope_qk_inplace(
query, key, cos_sin_cache, is_neox=False
)
else:
query, key = _apply_rotary_emb(
query, cos, sin, is_neox_style=False
), _apply_rotary_emb(key, cos, sin, is_neox_style=False)
attn_output = self.attn1(query, key, value, gate_compress=gate_compress) attn_output = self.attn1(query, key, value, gate_compress=gate_compress)
attn_output = attn_output.flatten(2) attn_output = attn_output.flatten(2)