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

apply_rotary_emb_vllm_flash_attn instance-attribute 

apply_rotary_emb_vllm_flash_attn = import_module(
    "vllm.vllm_flash_attn.layers.rotary"
).apply_rotary_emb

half_head_dim instance-attribute 

half_head_dim = half_head_dim

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

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

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)