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

[Diffusion] Refactor qwen_image's rope in a single helper func (#16047)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-12-29 17:24:26 +08:00
committed by GitHub
parent 4ab66d956f
commit f3d73b0199
5 changed files with 103 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py
View File

@@ -156,6 +156,7 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
# img_shapes: for global entire image
img_freqs, txt_freqs = rotary_emb(img_shapes, txt_seq_lens, device=device)
# flashinfer RoPE expects a float32 cos/sin cache concatenated on the last dim
img_cos_half = img_freqs.real.to(dtype=torch.float32).contiguous()
img_sin_half = img_freqs.imag.to(dtype=torch.float32).contiguous()
txt_cos_half = txt_freqs.real.to(dtype=torch.float32).contiguous()

68
python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
View File

@@ -27,7 +27,7 @@
"""Rotary Positional Embeddings."""
import functools
from collections import OrderedDict
from typing import Any
from typing import Any, Optional, Tuple
import torch
@@ -39,6 +39,72 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
def apply_flashinfer_rope_qk_inplace(
q: torch.Tensor,
k: torch.Tensor,
cos_sin_cache: torch.Tensor,
*,
head_size: Optional[int] = None,
is_neox: bool = False,
positions: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
if q.dim() != 4 or k.dim() != 4:
raise ValueError(
f"Expected q/k to be 4D [bsz, seqlen, nheads, head_size], "
f"got q:{tuple(q.shape)} k:{tuple(k.shape)}"
)
if q.shape != k.shape:
raise ValueError(
f"q and k must have the same shape, got {q.shape} vs {k.shape}"
)
if not (isinstance(cos_sin_cache, torch.Tensor) and cos_sin_cache.dim() == 2):
raise ValueError("cos_sin_cache must be a 2D torch.Tensor")
bsz, seqlen, nheads, d = q.shape
if head_size is None:
head_size = d
if head_size != d:
raise ValueError(f"head_size mismatch: inferred {d}, but head_size={head_size}")
try:
from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
except Exception as e:
raise RuntimeError(
"flashinfer is required for apply_flashinfer_rope_qk_inplace. "
"Please install flashinfer or disable this optimization."
) from e
if positions is None:
pos_1d = torch.arange(seqlen, device="cpu", dtype=torch.long)
positions = pos_1d if bsz == 1 else pos_1d.repeat(bsz)
else:
if not (
isinstance(positions, torch.Tensor)
and positions.dtype == torch.long
and positions.dim() == 1
):
raise ValueError("positions must be a 1D torch.long Tensor")
if positions.numel() != bsz * seqlen:
raise ValueError(
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(
positions=positions,
query=q_flat,
key=k_flat,
head_size=d,
cos_sin_cache=cos_sin_cache,
is_neox=is_neox,
)
return q_flat.view(bsz, seqlen, nheads, d), k_flat.view(bsz, seqlen, nheads, d)
def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :]

14
python/sglang/multimodal_gen/runtime/models/dits/flux.py
View File

@@ -40,7 +40,7 @@ from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
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.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import current_platform
@@ -182,11 +182,15 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
if freqs_cis is not None:
cos, sin = freqs_cis
query = _apply_rotary_emb(
query, cos, sin, is_neox_style=False, interleaved=False
cos_sin_cache = torch.cat(
[
cos.to(dtype=torch.float32).contiguous(),
sin.to(dtype=torch.float32).contiguous(),
],
dim=-1,
)
key = _apply_rotary_emb(
key, cos, sin, is_neox_style=False, interleaved=False
query, key = apply_flashinfer_rope_qk_inplace(
query, key, cos_sin_cache, is_neox=False
)
x = self.attn(query, key, value)

26
python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
View File

@@ -26,7 +26,7 @@ from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
NDRotaryEmbedding,
_apply_rotary_emb,
apply_flashinfer_rope_qk_inplace,
)
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import current_platform
@@ -187,11 +187,15 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
if freqs_cis is not None:
cos, sin = freqs_cis
query = _apply_rotary_emb(
query, cos, sin, is_neox_style=False, interleaved=True
cos_sin_cache = torch.cat(
[
cos.to(dtype=torch.float32).contiguous(),
sin.to(dtype=torch.float32).contiguous(),
],
dim=-1,
)
key = _apply_rotary_emb(
key, cos, sin, is_neox_style=False, interleaved=True
query, key = apply_flashinfer_rope_qk_inplace(
query, key, cos_sin_cache, is_neox=False
)
hidden_states = self.attn(query, key, value)
@@ -311,11 +315,15 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
if freqs_cis is not None:
cos, sin = freqs_cis
query = _apply_rotary_emb(
query, cos, sin, is_neox_style=False, interleaved=True
cos_sin_cache = torch.cat(
[
cos.to(dtype=torch.float32).contiguous(),
sin.to(dtype=torch.float32).contiguous(),
],
dim=-1,
)
key = _apply_rotary_emb(
key, cos, sin, is_neox_style=False, interleaved=True
query, key = apply_flashinfer_rope_qk_inplace(
query, key, cos_sin_cache, is_neox=False
)
hidden_states = self.attn(query, key, value)
hidden_states = hidden_states.flatten(2, 3)

53
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
View File

@@ -19,6 +19,9 @@ from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConf
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm, RMSNorm
from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
apply_flashinfer_rope_qk_inplace,
)
from sglang.multimodal_gen.runtime.layers.triton_ops import (
fuse_scale_shift_gate_select01_kernel,
fuse_scale_shift_kernel,
@@ -30,12 +33,6 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__) # pylint: disable=invalid-name
try:
from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
except Exception:
apply_rope_with_cos_sin_cache_inplace = None
def _get_qkv_projections(
attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None
):
@@ -561,9 +558,6 @@ class QwenImageCrossAttention(nn.Module):
# Apply RoPE
if image_rotary_emb is not None:
if apply_rope_with_cos_sin_cache_inplace is None:
raise RuntimeError("flashinfer is required")
if not (
isinstance(image_rotary_emb[0], torch.Tensor)
and image_rotary_emb[0].dim() == 2
@@ -572,41 +566,12 @@ class QwenImageCrossAttention(nn.Module):
img_cache, txt_cache = image_rotary_emb
def _apply_flashinfer_rope(
q_4d: torch.Tensor, k_4d: torch.Tensor, cache: torch.Tensor
):
bsz, seqlen, nheads, d = q_4d.shape
pos_1d = torch.arange(seqlen, device="cpu", dtype=torch.long)
if bsz == 1:
positions = pos_1d.to(q_4d.device, non_blocking=True)
q2 = q_4d.squeeze(0).reshape(seqlen, nheads * d).contiguous()
k2 = k_4d.squeeze(0).reshape(seqlen, nheads * d).contiguous()
apply_rope_with_cos_sin_cache_inplace(
positions=positions,
query=q2,
key=k2,
head_size=d,
cos_sin_cache=cache,
is_neox=False,
)
return q2.view(1, seqlen, nheads, d), k2.view(1, seqlen, nheads, d)
positions = pos_1d.repeat(bsz).to(q_4d.device, non_blocking=True)
q2 = q_4d.reshape(bsz * seqlen, nheads * d).contiguous()
k2 = k_4d.reshape(bsz * seqlen, nheads * d).contiguous()
apply_rope_with_cos_sin_cache_inplace(
positions=positions,
query=q2,
key=k2,
head_size=d,
cos_sin_cache=cache,
is_neox=False,
)
return q2.view(bsz, seqlen, nheads, d), k2.view(bsz, seqlen, nheads, d)
img_query, img_key = _apply_flashinfer_rope(img_query, img_key, img_cache)
txt_query, txt_key = _apply_flashinfer_rope(txt_query, txt_key, txt_cache)
img_query, img_key = apply_flashinfer_rope_qk_inplace(
img_query, img_key, img_cache, is_neox=False
)
txt_query, txt_key = apply_flashinfer_rope_qk_inplace(
txt_query, txt_key, txt_cache, is_neox=False
)
# Concatenate for joint attention
# Order: [text, image]