wehub-resource-sync
460 lines
14 KiB
Python
460 lines
14 KiB
Python
# Some code borrowed from the awesome work: https://github.com/zhuzilin/ring-flash-attention
|
|
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import torch
|
|
from functools import cache
|
|
|
|
from .ring_utils import RingComm
|
|
|
|
_NPU_BLOCK_MASK_SIZE = 2048
|
|
_NPU_FULL_TOKENS = 2147483647
|
|
_NPU_TND_SOFTMAX_STAT_REPEAT = 8
|
|
|
|
|
|
def is_npu_tensor(tensor: torch.Tensor) -> bool:
|
|
return tensor.device.type == 'npu'
|
|
|
|
|
|
def _cu_seqlens_to_actual_seq(cu_seqlens: torch.Tensor) -> tuple[int, ...]:
|
|
return tuple(int(x) for x in cu_seqlens[1:].detach().cpu().tolist())
|
|
|
|
|
|
@cache
|
|
def _get_npu_causal_mask_cpu() -> torch.Tensor:
|
|
return torch.triu(torch.ones((_NPU_BLOCK_MASK_SIZE, _NPU_BLOCK_MASK_SIZE), dtype=torch.bool), diagonal=1)
|
|
|
|
|
|
def _get_npu_causal_mask(device: torch.device) -> torch.Tensor:
|
|
return _get_npu_causal_mask_cpu().to(device=device)
|
|
|
|
|
|
def _normalize_window_size(window_size):
|
|
if window_size is None:
|
|
return -1, -1
|
|
return window_size
|
|
|
|
|
|
def _get_npu_sparse_params(causal: bool, window_size, device: torch.device) -> dict:
|
|
window_size = _normalize_window_size(window_size)
|
|
if window_size != (-1, -1):
|
|
left, right = window_size
|
|
left = _NPU_FULL_TOKENS if left < 0 else int(left)
|
|
right = _NPU_FULL_TOKENS if right < 0 else int(right)
|
|
if causal:
|
|
right = 0
|
|
return {
|
|
'atten_mask': _get_npu_causal_mask(device),
|
|
'sparse_mode': 4,
|
|
'pre_tockens': left,
|
|
'next_tockens': right,
|
|
}
|
|
if causal:
|
|
return {
|
|
'atten_mask': _get_npu_causal_mask(device),
|
|
'sparse_mode': 3,
|
|
'pre_tockens': _NPU_FULL_TOKENS,
|
|
'next_tockens': _NPU_FULL_TOKENS,
|
|
}
|
|
return {
|
|
'atten_mask': None,
|
|
'sparse_mode': 0,
|
|
'pre_tockens': _NPU_FULL_TOKENS,
|
|
'next_tockens': _NPU_FULL_TOKENS,
|
|
}
|
|
|
|
|
|
def _reshape_npu_lse(lse: torch.Tensor, seqlen_q: int, num_heads: int) -> torch.Tensor:
|
|
"""Normalize Ascend softmax stats to flash-attn's [num_heads, seqlen] layout."""
|
|
if lse.dim() == 2:
|
|
if lse.shape == (num_heads, seqlen_q):
|
|
return lse.contiguous()
|
|
if lse.shape == (seqlen_q, num_heads):
|
|
return lse.transpose(0, 1).contiguous()
|
|
elif lse.dim() == 3:
|
|
# Some CANN versions return an extra trailing size-8 axis with repeated
|
|
# stats. Ring merge only needs one copy of each token/head lse.
|
|
if lse.shape[-1] == 8:
|
|
lse = lse[..., 0]
|
|
if lse.shape == (seqlen_q, num_heads):
|
|
return lse.transpose(0, 1).contiguous()
|
|
if lse.shape == (num_heads, seqlen_q):
|
|
return lse.contiguous()
|
|
if lse.shape[0] == seqlen_q:
|
|
return lse.permute(1, 2, 0).reshape(num_heads, seqlen_q).contiguous()
|
|
if lse.shape[1] == seqlen_q:
|
|
return lse.permute(0, 2, 1).reshape(num_heads, seqlen_q).contiguous()
|
|
raise RuntimeError(f'Unexpected NPU lse shape {tuple(lse.shape)} for seqlen_q={seqlen_q}, num_heads={num_heads}')
|
|
|
|
|
|
def _get_npu_attention_common_kwargs(
|
|
q: torch.Tensor,
|
|
*,
|
|
cu_seqlens_q: torch.Tensor,
|
|
cu_seqlens_kv: torch.Tensor,
|
|
softmax_scale: float,
|
|
dropout_p: float,
|
|
causal: bool,
|
|
window_size,
|
|
deterministic: bool,
|
|
) -> dict:
|
|
sparse_params = _get_npu_sparse_params(causal, window_size, q.device)
|
|
return {
|
|
'head_num': q.shape[1],
|
|
'input_layout': 'TND',
|
|
'scale_value': softmax_scale or q.shape[-1]**(-0.5),
|
|
'keep_prob': 1. - dropout_p,
|
|
'actual_seq_qlen': _cu_seqlens_to_actual_seq(cu_seqlens_q),
|
|
'actual_seq_kvlen': _cu_seqlens_to_actual_seq(cu_seqlens_kv),
|
|
'sync': bool(deterministic and dropout_p > 0),
|
|
**sparse_params,
|
|
}
|
|
|
|
|
|
def _call_npu_fusion_attention(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
*,
|
|
cu_seqlens_q: torch.Tensor,
|
|
cu_seqlens_kv: torch.Tensor,
|
|
softmax_scale: float,
|
|
dropout_p: float,
|
|
causal: bool,
|
|
window_size,
|
|
deterministic: bool,
|
|
):
|
|
import torch_npu
|
|
|
|
common_kwargs = _get_npu_attention_common_kwargs(
|
|
q,
|
|
cu_seqlens_q=cu_seqlens_q,
|
|
cu_seqlens_kv=cu_seqlens_kv,
|
|
softmax_scale=softmax_scale,
|
|
dropout_p=dropout_p,
|
|
causal=causal,
|
|
window_size=window_size,
|
|
deterministic=deterministic,
|
|
)
|
|
params = {
|
|
'query': q,
|
|
'key': k,
|
|
'value': v,
|
|
'scale': common_kwargs['scale_value'],
|
|
'softmax_layout': 'TND',
|
|
}
|
|
params.update(common_kwargs)
|
|
params.pop('scale_value')
|
|
return torch_npu.npu_fusion_attention(**params)
|
|
|
|
|
|
def _call_npu_fusion_attention_grad(
|
|
dout: torch.Tensor,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
*,
|
|
attention_out: torch.Tensor,
|
|
softmax_lse: torch.Tensor,
|
|
cu_seqlens_q: torch.Tensor,
|
|
cu_seqlens_kv: torch.Tensor,
|
|
softmax_scale: float,
|
|
dropout_p: float,
|
|
causal: bool,
|
|
window_size,
|
|
deterministic: bool,
|
|
):
|
|
import torch_npu
|
|
|
|
if not hasattr(torch_npu, 'npu_fusion_attention_grad'):
|
|
raise AttributeError('torch_npu.npu_fusion_attention_grad is not available')
|
|
# Dropout backward needs the exact seed/offset from the original forward,
|
|
# which this ring ctx does not save. Fail instead of using a wrong mask.
|
|
if dropout_p != 0.0:
|
|
raise NotImplementedError('NPU ring attention native backward currently requires dropout_p=0.')
|
|
|
|
common_kwargs = _get_npu_attention_common_kwargs(
|
|
q,
|
|
cu_seqlens_q=cu_seqlens_q,
|
|
cu_seqlens_kv=cu_seqlens_kv,
|
|
softmax_scale=softmax_scale,
|
|
dropout_p=dropout_p,
|
|
causal=causal,
|
|
window_size=window_size,
|
|
deterministic=deterministic,
|
|
)
|
|
softmax_max, softmax_sum = _npu_softmax_stats_from_global_lse(
|
|
softmax_lse,
|
|
q_tokens=q.shape[0],
|
|
num_heads=q.shape[1],
|
|
)
|
|
|
|
params = {
|
|
'query': q,
|
|
'key': k,
|
|
'value': v,
|
|
'dy': dout,
|
|
'head_num': common_kwargs['head_num'],
|
|
'input_layout': common_kwargs['input_layout'],
|
|
'atten_mask': common_kwargs['atten_mask'],
|
|
'softmax_max': softmax_max,
|
|
'softmax_sum': softmax_sum,
|
|
'softmax_in': None,
|
|
'attention_in': (attention_out if torch.is_tensor(attention_out) and attention_out.numel() > 0 else None),
|
|
'scale_value': common_kwargs['scale_value'],
|
|
'keep_prob': common_kwargs['keep_prob'],
|
|
'pre_tockens': common_kwargs['pre_tockens'],
|
|
'next_tockens': common_kwargs['next_tockens'],
|
|
'seed': 0,
|
|
'offset': 0,
|
|
'numels': 0,
|
|
'actual_seq_qlen': common_kwargs['actual_seq_qlen'],
|
|
'actual_seq_kvlen': common_kwargs['actual_seq_kvlen'],
|
|
'sparse_mode': common_kwargs['sparse_mode'],
|
|
'sync': common_kwargs['sync'],
|
|
'softmax_layout': 'TND',
|
|
}
|
|
return torch_npu.npu_fusion_attention_grad(**params)
|
|
|
|
|
|
def _normalize_flash_attn_lse(softmax_lse: torch.Tensor, total_len: int) -> torch.Tensor:
|
|
"""Normalize flash-attn lse to [num_heads, total_len]."""
|
|
lse = softmax_lse
|
|
if lse.dim() == 3 and lse.shape[0] == 1:
|
|
lse = lse.squeeze(0)
|
|
if lse.dim() != 2:
|
|
raise RuntimeError(f'Unexpected softmax_lse shape: {tuple(softmax_lse.shape)}')
|
|
if lse.shape[1] != total_len:
|
|
lse = lse.transpose(0, 1).contiguous()
|
|
if lse.shape[1] != total_len:
|
|
raise RuntimeError(f'Unexpected softmax_lse shape: {tuple(softmax_lse.shape)} for total_len={total_len}')
|
|
return lse
|
|
|
|
|
|
def _npu_softmax_stats_from_global_lse(
|
|
softmax_lse_global: torch.Tensor,
|
|
q_tokens: int,
|
|
num_heads: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
lse_h_t = _normalize_flash_attn_lse(softmax_lse_global, q_tokens)
|
|
if lse_h_t.shape[0] != num_heads:
|
|
raise RuntimeError(f'Unexpected global lse shape: {tuple(softmax_lse_global.shape)} '
|
|
f'for q_tokens={q_tokens}, num_heads={num_heads}')
|
|
|
|
# With softmax_layout='TND', Ascend returns softmax stats as [T, N, 8].
|
|
# The split-attention backward only needs logsumexp; max=lse and sum=1
|
|
# encode the same value without replaying the block forward.
|
|
lse_t_h = lse_h_t.transpose(0, 1).contiguous().to(torch.float32)
|
|
softmax_max = lse_t_h.unsqueeze(-1).expand(
|
|
q_tokens,
|
|
num_heads,
|
|
_NPU_TND_SOFTMAX_STAT_REPEAT,
|
|
).contiguous()
|
|
return softmax_max, torch.ones_like(softmax_max)
|
|
|
|
|
|
def _get_second_half_lse(softmax_lse: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
|
|
total_len = int(cu_seqlens[-1].item())
|
|
lse = _normalize_flash_attn_lse(softmax_lse, total_len)
|
|
|
|
# The step > rank branch only differentiates q[half_index1]. Slice the final
|
|
# merged lse per sequence so the native grad ctx sees the same query span.
|
|
second_half_lse = torch.empty((lse.shape[0], lse.shape[1] // 2), dtype=lse.dtype, device=lse.device)
|
|
for i in range(len(cu_seqlens) - 1):
|
|
start, end = cu_seqlens[i].item(), cu_seqlens[i + 1].item()
|
|
new_start, new_end = start // 2, end // 2
|
|
start += (end - start) // 2
|
|
second_half_lse[:, new_start:new_end] = lse[:, start:end]
|
|
return second_half_lse
|
|
|
|
|
|
def _npu_block_backward_with_global_stats(
|
|
block_dout,
|
|
block_q,
|
|
block_k,
|
|
block_v,
|
|
block_out_global,
|
|
block_lse_global,
|
|
block_causal,
|
|
cu_seqlens_q,
|
|
cu_seqlens_kv,
|
|
softmax_scale,
|
|
dropout_p,
|
|
window_size,
|
|
deterministic,
|
|
):
|
|
"""Run one native NPU block backward using the final merged ring stats."""
|
|
return _call_npu_fusion_attention_grad(
|
|
block_dout.to(block_q.dtype),
|
|
block_q,
|
|
block_k,
|
|
block_v,
|
|
attention_out=block_out_global,
|
|
softmax_lse=block_lse_global,
|
|
cu_seqlens_q=cu_seqlens_q,
|
|
cu_seqlens_kv=cu_seqlens_kv,
|
|
softmax_scale=softmax_scale,
|
|
dropout_p=dropout_p,
|
|
causal=block_causal,
|
|
window_size=window_size,
|
|
deterministic=deterministic,
|
|
)[:3]
|
|
|
|
|
|
def _squeeze_batch(*tensors):
|
|
squeezed = []
|
|
for tensor in tensors:
|
|
if tensor.shape[0] == 1:
|
|
squeezed.append(tensor.squeeze(0))
|
|
else:
|
|
squeezed.append(tensor)
|
|
return tuple(squeezed)
|
|
|
|
|
|
def npu_backward(
|
|
process_group,
|
|
dout,
|
|
q,
|
|
k,
|
|
v,
|
|
out,
|
|
softmax_lse,
|
|
cu_seqlens,
|
|
max_seqlen,
|
|
half_index0,
|
|
half_index1,
|
|
softmax_scale,
|
|
dropout_p=0.0,
|
|
window_size=(-1, -1),
|
|
deterministic=False,
|
|
):
|
|
kv_comm = RingComm(process_group)
|
|
d_kv_comm = RingComm(process_group)
|
|
dout, q, k, v, out, softmax_lse = _squeeze_batch(dout, q, k, v, out, softmax_lse)
|
|
cu_seqlens = cu_seqlens // kv_comm.world_size
|
|
del max_seqlen
|
|
|
|
half_cu_seqlens = cu_seqlens // 2
|
|
q1 = q[half_index1]
|
|
dout1 = dout[half_index1]
|
|
out1 = out[half_index1]
|
|
softmax_lse1 = _get_second_half_lse(softmax_lse, cu_seqlens)
|
|
|
|
dq = torch.zeros_like(q, dtype=torch.float32)
|
|
current_step_dk = torch.empty_like(k, dtype=torch.float32)
|
|
current_step_dv = torch.empty_like(v, dtype=torch.float32)
|
|
next_dk = next_dv = None
|
|
|
|
for step in range(kv_comm.world_size):
|
|
current_step_dk.zero_()
|
|
current_step_dv.zero_()
|
|
if step == 0:
|
|
bdq, bdk, bdv = _npu_block_backward_with_global_stats(
|
|
dout,
|
|
q,
|
|
k,
|
|
v,
|
|
out,
|
|
softmax_lse,
|
|
True,
|
|
cu_seqlens,
|
|
cu_seqlens,
|
|
softmax_scale,
|
|
dropout_p,
|
|
window_size,
|
|
deterministic,
|
|
)
|
|
dq += bdq.to(torch.float32)
|
|
current_step_dk += bdk.to(torch.float32)
|
|
current_step_dv += bdv.to(torch.float32)
|
|
elif step <= kv_comm.rank:
|
|
k0 = k[half_index0]
|
|
v0 = v[half_index0]
|
|
bdq, bdk, bdv = _npu_block_backward_with_global_stats(
|
|
dout,
|
|
q,
|
|
k0,
|
|
v0,
|
|
out,
|
|
softmax_lse,
|
|
False,
|
|
cu_seqlens,
|
|
half_cu_seqlens,
|
|
softmax_scale,
|
|
dropout_p,
|
|
window_size,
|
|
deterministic,
|
|
)
|
|
dq += bdq.to(torch.float32)
|
|
current_step_dk[half_index0] += bdk.to(torch.float32)
|
|
current_step_dv[half_index0] += bdv.to(torch.float32)
|
|
else:
|
|
bdq, bdk, bdv = _npu_block_backward_with_global_stats(
|
|
dout1,
|
|
q1,
|
|
k,
|
|
v,
|
|
out1,
|
|
softmax_lse1,
|
|
False,
|
|
half_cu_seqlens,
|
|
cu_seqlens,
|
|
softmax_scale,
|
|
dropout_p,
|
|
window_size,
|
|
deterministic,
|
|
)
|
|
dq[half_index1] += bdq.to(torch.float32)
|
|
current_step_dk += bdk.to(torch.float32)
|
|
current_step_dv += bdv.to(torch.float32)
|
|
|
|
# K/V gradients are owned by the rank that originally held that shard.
|
|
# Rotate the accumulated gradients in the opposite ring direction until
|
|
# each owner receives its final dk/dv.
|
|
if step == 0:
|
|
dk = current_step_dk
|
|
dv = current_step_dv
|
|
else:
|
|
dk = next_dk
|
|
dv = next_dv
|
|
dk += current_step_dk
|
|
dv += current_step_dv
|
|
|
|
next_dk, next_dv = d_kv_comm.send_recv_kv(dk, dv)
|
|
d_kv_comm.wait()
|
|
|
|
if step + 1 != kv_comm.world_size:
|
|
next_k, next_v = kv_comm.send_recv_kv(k, v)
|
|
kv_comm.wait()
|
|
k, v = next_k, next_v
|
|
|
|
return dq.to(q.dtype).unsqueeze(0), next_dk.to(q.dtype).unsqueeze(0), next_dv.to(q.dtype).unsqueeze(0)
|
|
|
|
|
|
def npu_forward(
|
|
q,
|
|
k,
|
|
v,
|
|
causal,
|
|
cu_seqlens_q,
|
|
cu_seqlens_kv,
|
|
dropout_p,
|
|
softmax_scale,
|
|
deterministic,
|
|
window_size,
|
|
):
|
|
outputs = _call_npu_fusion_attention(
|
|
q,
|
|
k,
|
|
v,
|
|
cu_seqlens_q=cu_seqlens_q,
|
|
cu_seqlens_kv=cu_seqlens_kv,
|
|
softmax_scale=softmax_scale,
|
|
dropout_p=dropout_p,
|
|
causal=causal,
|
|
window_size=window_size,
|
|
deterministic=deterministic,
|
|
)
|
|
block_out, softmax_max, softmax_sum = outputs[:3]
|
|
block_lse = softmax_max.to(torch.float32) + torch.log(softmax_sum.to(torch.float32))
|
|
block_lse = _reshape_npu_lse(block_lse, q.shape[0], q.shape[1])
|
|
return block_out, block_lse
|