[NPU]Replace ascend vision attn operator (#25768)
Signed-off-by: zhaozx-cn <zhaozx2116@163.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
github-actions[bot]
parent
0ae4740bd1
commit
224b1dc775
@@ -733,27 +733,26 @@ class VisionAscendAttention(nn.Module):
|
|||||||
seq_len_arg = cu_seqlens
|
seq_len_arg = cu_seqlens
|
||||||
else:
|
else:
|
||||||
cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device="cpu")
|
cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device="cpu")
|
||||||
seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
|
seq_len_arg = cu_seqlens[1:].to(torch.int32)
|
||||||
if seq_lens.is_npu:
|
|
||||||
seq_lens = seq_lens.to("cpu")
|
|
||||||
output = torch.empty_like(q)
|
|
||||||
seq_len_arg = seq_lens.to(torch.int32)
|
|
||||||
|
|
||||||
_, num_heads, head_size = q.shape
|
_, num_heads, head_size = q.shape
|
||||||
num_kv_heads = k.shape[1]
|
num_kv_heads = k.shape[1]
|
||||||
|
|
||||||
scale_value = softmax_scale if softmax_scale is not None else head_size**-0.5
|
scale_value = softmax_scale if softmax_scale is not None else head_size**-0.5
|
||||||
|
|
||||||
torch_npu._npu_flash_attention_unpad(
|
seq_len_arg = seq_len_arg.tolist()
|
||||||
|
output = torch_npu.npu_fused_infer_attention_score(
|
||||||
query=q,
|
query=q,
|
||||||
key=k,
|
key=k,
|
||||||
value=v,
|
value=v,
|
||||||
seq_len=seq_len_arg,
|
actual_seq_lengths=seq_len_arg,
|
||||||
scale_value=scale_value,
|
actual_seq_lengths_kv=seq_len_arg,
|
||||||
|
scale=scale_value,
|
||||||
num_heads=num_heads,
|
num_heads=num_heads,
|
||||||
num_kv_heads=num_kv_heads,
|
num_key_value_heads=num_kv_heads,
|
||||||
out=output,
|
sparse_mode=0,
|
||||||
)
|
input_layout="TND",
|
||||||
|
)[0]
|
||||||
return output
|
return output
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user