From 294ff71d189c918205e29181233c37121e1a5056 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Thu, 8 Jan 2026 22:44:38 +0800 Subject: [PATCH] [Diffusion] Avoid cpu2gpu sync in flashinfer rope and apply flashinfer rope to wanvideo (#16668) Co-authored-by: Mick --- .../runtime/layers/rotary_embedding.py | 4 +- .../runtime/models/dits/wanvideo.py | 37 ++++++++++++++++--- 2 files changed, 32 insertions(+), 9 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py index 2ef943229..1bf22fcf7 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py @@ -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( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 97ada9791..ee63a77ce 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -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)