vllm_omni.platforms.npu.layers.rotary_embedding ¶

Shared NPU rotary embedding helpers.

npu_rotary_mul_with_bsnd_fallback ¶

npu_rotary_mul_with_bsnd_fallback(
    hidden_states: Tensor,
    cos: Tensor,
    sin: Tensor,
    unsqueeze_dim: int,
) -> Tensor

Apply fused rotary multiplication, falling back from BNSD to BSND.

unsqueeze_dim=1 denotes BNSD input following the Transformers rotary embedding convention. Supported BNSD shapes stay on the fast path. Shapes rejected by the CANN BNSD tiler are transposed to BSND for the fused op and transposed back before returning.