Skip to content

vllm_omni.diffusion.models.wan2_2.rope

Wan S2V rotary positional embedding layers.

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