chenchenghao/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
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

4
python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
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)
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)
else:
if not (
@@ -108,8 +108,6 @@ def apply_flashinfer_rope_qk_inplace(
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()
k_flat = k.reshape(bsz * seqlen, nheads * d).contiguous()
apply_rope_with_cos_sin_cache_inplace(

37
python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
View File

@@ -30,6 +30,7 @@ from sglang.multimodal_gen.runtime.layers.mlp import MLP
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
NDRotaryEmbedding,
_apply_rotary_emb,
apply_flashinfer_rope_qk_inplace,
)
from sglang.multimodal_gen.runtime.layers.visual_embedding import (
ModulateProjection,
@@ -399,9 +400,21 @@ class WanTransformerBlock(nn.Module):
# Apply rotary embeddings
cos, sin = freqs_cis
query, key = _apply_rotary_emb(
query, cos, sin, is_neox_style=False
), _apply_rotary_emb(key, cos, sin, is_neox_style=False)
if query.is_cuda and query.shape == key.shape:
cos_sin_cache = torch.cat(
[
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 = attn_output.flatten(2)
attn_output, _ = self.to_out(attn_output)
@@ -569,9 +582,21 @@ class WanTransformerBlock_VSA(nn.Module):
# Apply rotary embeddings
cos, sin = freqs_cis
query, key = _apply_rotary_emb(
query, cos, sin, is_neox_style=False
), _apply_rotary_emb(key, cos, sin, is_neox_style=False)
if query.is_cuda and query.shape == key.shape:
cos_sin_cache = torch.cat(
[
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 = attn_output.flatten(2)