Skip to content

vllm_omni.diffusion.layers.rope

logger module-attribute 

logger = init_logger(__name__)

RotaryEmbedding

Bases: CustomOp

rotary positional embedding. interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of 1st half and 2nd half (GPT-NeoX style).

apply_rotary_emb_flash_attn instance-attribute 

apply_rotary_emb_flash_attn = apply_rotary

has_mindie instance-attribute 

has_mindie = False

interleaved instance-attribute 

interleaved = not is_neox_style

is_neox_style instance-attribute 

is_neox_style = is_neox_style

forward_cuda 

forward_cuda(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor

forward_hip 

forward_hip(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor

forward_musa 

forward_musa(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor

forward_native 

forward_native(
    x: Tensor, cos: Tensor, sin: Tensor
) -> Tensor

forward_npu 

forward_npu(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor

forward_xpu 

forward_xpu(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor

RotaryEmbeddingS2VGrid

Bases: Module

Grid-based RoPE for S2V motioner/init attention.

Applies complex-valued rotary embeddings using 3D grid sampling (frame, height, width). Used by SimpleSelfAttention, SwinSelfAttention, CausalSelfAttention in motioner blocks.

apply_precomputed staticmethod 

apply_precomputed(
    x: Tensor, precomputed_freqs: Tensor
) -> Tensor

Apply precomputed position frequencies to input tensor.

forward staticmethod 

forward(
    x: Tensor, grid_sizes, freqs: Tensor, start=None
) -> Tensor

precompute staticmethod 

precompute(
    seq_len: int,
    num_heads: int,
    head_dim: int,
    grid_sizes,
    freqs: Tensor,
    device: device,
    start=None,
) -> Tensor

Precompute position frequency tensor from grid specification.

Returns a complex tensor that can be reused across layers via apply_precomputed().

RotaryEmbeddingWan

Bases: RotaryEmbedding

rotary positional embedding for Wan. interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of 1st half and 2nd half (GPT-NeoX style).

half_head_dim instance-attribute 

half_head_dim = half_head_dim

forward_cuda 

forward_cuda(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor

forward_hip 

forward_hip(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor

forward_native 

forward_native(
    x: Tensor, cos: Tensor, sin: Tensor
) -> Tensor

forward_npu 

forward_npu(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor

RotaryEmbeddingWanS2V

Bases: RotaryEmbeddingWan

Apply RoPE using precomputed complex freqs for Wan S2V main transformer.

Converts complex freqs (from WanS2VRotaryPosEmbed) to cos/sin and delegates to RotaryEmbeddingWan for platform-optimized application (float32 kernel). Under TP, freqs has 1 head — broadcasts automatically via cos/sin.

forward 

forward(x: Tensor, freqs: Tensor) -> Tensor

WanS2VRotaryPosEmbed

Bases: Module

Precompute complex-valued RoPE embeddings for S2V multi-grid positions.

Owns the base frequency buffer and provides forward() to compute position embeddings given hidden_states (for shape) and grid_sizes.

head_dim instance-attribute 

head_dim = head_dim

num_heads instance-attribute 

num_heads = num_heads

forward 

forward(
    hidden_states: Tensor,
    grid_sizes: list,
    trainable_freqs: Tensor | None = None,
) -> Tensor

Precompute RoPE embeddings for the given grid layout.

Parameters:

Name Type Description Default
hidden_states Tensor

Tensor [B, S, ...] (used for batch/seq shape and device)

 required
grid_sizes list

Grid specification (list of [offsets, sizes, totals])

 required
trainable_freqs Tensor | None

Optional trainable frequency overrides for t_f < 0

 None

Returns:

Type Description
Tensor

Complex tensor [B, S, 1, head_dim//2] of precomputed position embeddings

apply_rope_to_qk 

apply_rope_to_qk(
    rope: RotaryEmbedding,
    query: Tensor,
    key: Tensor,
    image_rotary_emb: tuple[Tensor, Tensor] | None,
) -> tuple[Tensor, Tensor]

Apply rotary positional embeddings to query and key tensors.

Parameters:

Name Type Description Default
rope RotaryEmbedding

RotaryEmbedding instance for applying position embeddings

 required
query Tensor

Query tensor [B, S, H, D]

 required
key Tensor

Key tensor [B, S, H, D]

 required
image_rotary_emb tuple[Tensor, Tensor] | None

Tuple of (cos, sin) tensors or None

 required

Returns:

Type Description
tuple[Tensor, Tensor]

Tuple of (query, key) with RoPE applied if rotary embeddings provided

apply_rotary_emb_mindiesd 

apply_rotary_emb_mindiesd(
    x: Tensor,
    cos: Tensor,
    sin: Tensor,
    interleaved: bool = False,
    half_head_dim: bool = True,
) -> Tensor

apply_rotary_emb_torch 

apply_rotary_emb_torch(x, cos, sin, interleaved=False)

x: (batch_size, seqlen, nheads, headdim) cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)

rotate_half 

rotate_half(x, interleaved=False)