58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Shared NPU rotary embedding helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch_npu
|
|
|
|
|
|
def _bnsd_rotary_shape_is_supported(hidden_states: torch.Tensor) -> bool:
|
|
"""Return whether the fused half-mode BNSD tiler supports the shape.
|
|
|
|
For the aligned-head-dimension branch, CANN requires
|
|
``batch * num_heads <= seq_len * 8``. Other alignments are handled with
|
|
the BSND layout, which does not have this short-sequence limit.
|
|
"""
|
|
if hidden_states.ndim != 4:
|
|
return False
|
|
|
|
batch_size, num_heads, seq_len, head_dim = hidden_states.shape
|
|
element_size = hidden_states.element_size()
|
|
if head_dim % 2 != 0 or element_size <= 0 or 32 % element_size != 0:
|
|
return False
|
|
|
|
half_dim_alignment = 32 // element_size
|
|
return head_dim // 2 % half_dim_alignment == 0 and batch_size * num_heads <= seq_len * 8
|
|
|
|
|
|
def npu_rotary_mul_with_bsnd_fallback(
|
|
hidden_states: torch.Tensor,
|
|
cos: torch.Tensor,
|
|
sin: torch.Tensor,
|
|
unsqueeze_dim: int,
|
|
) -> torch.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.
|
|
"""
|
|
if unsqueeze_dim != 1 or _bnsd_rotary_shape_is_supported(hidden_states):
|
|
return torch_npu.npu_rotary_mul(
|
|
hidden_states,
|
|
cos.unsqueeze(unsqueeze_dim),
|
|
sin.unsqueeze(unsqueeze_dim),
|
|
)
|
|
|
|
hidden_states_bsnd = hidden_states.transpose(1, 2).contiguous()
|
|
output_bsnd = torch_npu.npu_rotary_mul(
|
|
hidden_states_bsnd,
|
|
cos.unsqueeze(2),
|
|
sin.unsqueeze(2),
|
|
)
|
|
return output_bsnd.transpose(1, 2)
|