vllm_omni.diffusion.layers.indexed_modulation ¶

Fused indexed modulation with FP32 accumulation.

indexed_gate ¶

indexed_gate(
    x: Tensor, gate: Tensor, other: Tensor, indices: Tensor
) -> Tensor

Return x + gate[indices] * other without indexed temporaries.

indexed_gate_rms_norm_scale_shift ¶

indexed_gate_rms_norm_scale_shift(
    residual: Tensor,
    gate: Tensor,
    branch: Tensor,
    weight: Tensor,
    shift: Tensor,
    scale: Tensor,
    indices: Tensor,
    eps: float,
) -> tuple[Tensor, Tensor]

Fuse gated residual, the following RMSNorm, and AdaLN affine.

indexed_scale_shift_ ¶

indexed_scale_shift_(
    x: Tensor, shift: Tensor, scale: Tensor, indices: Tensor
) -> Tensor

Apply indexed scale/shift in-place to a disposable contiguous input.

rms_norm_indexed_scale_shift ¶

rms_norm_indexed_scale_shift(
    x: Tensor,
    weight: Tensor,
    shift: Tensor,
    scale: Tensor,
    indices: Tensor,
    eps: float,
) -> Tensor

Fuse RMSNorm with an indexed affine transform.