mirror of https://github.com/hpcaitech/ColossalAI
314 lines
13 KiB
Python
314 lines
13 KiB
Python
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from transformers.modeling_attn_mask_utils import AttentionMaskConverter
|
|
|
|
|
|
def copy_to_cache(source, cache, lengths, block_tables, type: str = "prefill"):
|
|
"""
|
|
Func: copy key/value into key/value cache.
|
|
|
|
Args: key/value(source): shape [bsz,seq_len,num_heads,head_size]
|
|
cache: shape [num_blocks, num_kv_heads, head_size, block_size]
|
|
lengths: key/value lengths
|
|
block_tables
|
|
"""
|
|
num_blocks, num_heads, block_size, head_size = cache.shape
|
|
bsz, max_blocks_per_seq = block_tables.shape
|
|
needed_blocks = (lengths + block_size - 1) // block_size
|
|
|
|
if type == "prefill":
|
|
for i in range(bsz):
|
|
seq_len = lengths[i]
|
|
block_num = needed_blocks[i]
|
|
token_id = 0
|
|
for block_idx in range(block_num - 1):
|
|
cache[block_tables[i][block_idx]] = source[i][token_id : token_id + block_size].permute(1, 0, 2)
|
|
token_id += block_size
|
|
cache[block_tables[i][block_num - 1], :, : seq_len - token_id, :] = source[i][token_id:seq_len].permute(
|
|
1, 0, 2
|
|
)
|
|
elif type == "decoding":
|
|
assert source.size(1) == 1, "seq_len should be equal to 1 when decoding."
|
|
source = source.squeeze(1)
|
|
slot_idx = (lengths + block_size - 1) % block_size
|
|
for i in range(bsz):
|
|
cache[block_tables[i, needed_blocks[i] - 1], :, slot_idx[i], :] = source[i]
|
|
|
|
return cache
|
|
|
|
|
|
def convert_kvcache(cache, lengths, block_tables, pad_id=0):
|
|
"""
|
|
Func: convert key/value cache for calculation
|
|
|
|
Args: cache: shape [num_blocks, num_heads, block_size, head_size]
|
|
lengths: key/value length
|
|
block_tables
|
|
pad_id: padded_id
|
|
"""
|
|
num_blocks, num_heads, block_size, head_size = cache.shape
|
|
|
|
needed_blocks = (lengths + block_size - 1) // block_size
|
|
num_remaing_tokens = lengths % block_size
|
|
num_remaing_tokens[num_remaing_tokens == 0] += block_size
|
|
bsz = block_tables.shape[0]
|
|
seq_len = max(lengths)
|
|
padded_cache = []
|
|
for i in range(bsz):
|
|
_cache = torch.cat(
|
|
(
|
|
cache[block_tables[i][: needed_blocks[i] - 1]].permute((0, 2, 1, 3)).reshape(-1, num_heads, head_size),
|
|
cache[block_tables[i][needed_blocks[i] - 1], :, : num_remaing_tokens[i], :].permute(1, 0, 2),
|
|
),
|
|
dim=0,
|
|
)
|
|
padding = seq_len - _cache.size(0)
|
|
if padding > 0:
|
|
_cache = F.pad(_cache, (0, 0, 0, 0, 0, padding), value=pad_id)
|
|
padded_cache.append(_cache)
|
|
return torch.stack(padded_cache, dim=0)
|
|
|
|
|
|
class PagedAttention:
|
|
"""
|
|
Pure Torch implementation version of paged_attention.
|
|
Holds different types of forward function and useful components.
|
|
"""
|
|
|
|
@staticmethod
|
|
def pad_and_reshape(tensor, seq_lengths, max_seq_len, num_heads, head_size):
|
|
"""
|
|
Transform 1D no_pad tensor into 2D padded tensor with shape [bsz,seq_len,num_heads,head_size]
|
|
"""
|
|
bsz = len(seq_lengths)
|
|
padded_tensor = torch.zeros(bsz, max_seq_len, num_heads, head_size, dtype=tensor.dtype)
|
|
|
|
token_idx = 0
|
|
for i, seq_len in enumerate(seq_lengths):
|
|
seq_tensor = tensor[token_idx : token_idx + seq_len]
|
|
padded_tensor[i, :seq_len, :, :] = seq_tensor
|
|
token_idx += seq_len
|
|
return padded_tensor
|
|
|
|
@staticmethod
|
|
def generate_padding_mask(lengths, max_seq_len):
|
|
range_tensor = torch.arange(max_seq_len).expand(len(lengths), max_seq_len)
|
|
padding_mask = range_tensor < lengths.unsqueeze(1)
|
|
return padding_mask
|
|
|
|
@staticmethod
|
|
def repeat_kv(hidden_states: torch.Tensor, n_rep: int = 1) -> torch.Tensor:
|
|
"""
|
|
Essential component for MQA. Equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep).
|
|
Args: hidden_states(batch, num_key_value_heads, seqlen, head_dim)
|
|
n_rep: times of repeatition.
|
|
Output: hidden_states (batch, num_attention_heads, seqlen, head_dim)
|
|
"""
|
|
if n_rep == 1:
|
|
return hidden_states
|
|
|
|
batch, num_key_value_heads, seq_len, head_dim = hidden_states.shape
|
|
num_attention_heads = n_rep * num_key_value_heads
|
|
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, seq_len, head_dim)
|
|
|
|
return hidden_states.reshape(batch, num_attention_heads, seq_len, head_dim)
|
|
|
|
@staticmethod
|
|
def nopad_context_forward(
|
|
q: torch.Tensor, # [num_tokens, num_heads, head_size]
|
|
k: torch.Tensor, # [num_tokens, num_kv_heads, head_size]
|
|
v: torch.Tensor,
|
|
k_cache: torch.Tensor, # [num_blocks, num_heads, block_size, head_size]
|
|
v_cache: torch.Tensor,
|
|
context_lengths: torch.Tensor, # [num_seqs]
|
|
block_tables: torch.Tensor, # [num_seqs,max_blocks_per_sequence]
|
|
):
|
|
"""
|
|
NOTE: q,k,v are projected and applied rotary embedding, all aligned with triton version.
|
|
"""
|
|
# Fisrt, do shape verification
|
|
num_tokens, num_heads, head_size = q.shape
|
|
num_kv_heads = k.shape[-2]
|
|
|
|
assert num_heads % num_kv_heads == 0, "num_kv_heads should be divisible by num_heads"
|
|
num_kv_groups = num_heads // num_kv_heads
|
|
|
|
block_size = k_cache.size(-2)
|
|
bsz, max_blocks_per_sequence = block_tables.shape
|
|
max_seq_len = max_blocks_per_sequence * block_size
|
|
assert q.shape[-1] == k.shape[-1] == v.shape[-1]
|
|
assert q.shape[0] == k.shape[0] == v.shape[0]
|
|
assert context_lengths.shape[0] == block_tables.shape[0]
|
|
shape = (bsz, max_seq_len, num_heads, head_size)
|
|
input_shape = shape[:2]
|
|
|
|
q = PagedAttention.pad_and_reshape(
|
|
q, context_lengths, max_seq_len, num_heads, head_size
|
|
) # bsz,seqlen,num_heads,head_size
|
|
k = PagedAttention.pad_and_reshape(k, context_lengths, max_seq_len, num_heads, head_size)
|
|
v = PagedAttention.pad_and_reshape(v, context_lengths, max_seq_len, num_heads, head_size)
|
|
|
|
copy_to_cache(k, k_cache, lengths=context_lengths, block_tables=block_tables)
|
|
copy_to_cache(v, v_cache, lengths=context_lengths, block_tables=block_tables)
|
|
|
|
attn_mask = AttentionMaskConverter._make_causal_mask(input_shape, q.dtype, q.device, past_key_values_length=0)
|
|
attn_mask = attn_mask + PagedAttention.generate_padding_mask(context_lengths, max_seq_len)
|
|
|
|
q = q.transpose(1, 2)
|
|
k = PagedAttention.repeat_kv(k.transpose(1, 2), num_kv_groups)
|
|
v = PagedAttention.repeat_kv(v.transpose(1, 2), num_kv_groups)
|
|
|
|
# position_ids = torch.arange(0, max_seq_len, dtype=torch.long, device=query.device)
|
|
# position_ids = position_ids.unsqueeze(0)
|
|
# cos, sin = self.rotary_emb(value, max_seq_len)
|
|
# query, key = apply_rotary_pos_emb(query, key, cos, sin, position_ids)
|
|
|
|
attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(head_size)
|
|
if attn_weights.size() != (bsz, num_heads, max_seq_len, max_seq_len):
|
|
raise ValueError(f"Got wrong attn_weights, should be in shape {(bsz,num_heads,max_seq_len,max_seq_len)}.")
|
|
|
|
if attn_mask is not None:
|
|
attn_weights += attn_mask
|
|
|
|
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
|
|
attn_output = torch.matmul(attn_weights, v)
|
|
|
|
if attn_output.size() != (bsz, num_heads, max_seq_len, head_size):
|
|
raise ValueError(f"Got wrong attn_output, should be in shape {(bsz,num_heads,max_seq_len,head_size)}.")
|
|
attn_output = attn_output.transpose(1, 2).contiguous().reshape(bsz, max_seq_len, -1)
|
|
|
|
del attn_weights
|
|
|
|
return attn_output
|
|
|
|
@staticmethod
|
|
def pad_context_forward(
|
|
q: torch.Tensor, # [batch_size, seq_len, num_heads, head_size]
|
|
k: torch.Tensor, # [batch_size, seq_len, num_kv_heads, head_size]
|
|
v: torch.Tensor,
|
|
k_cache: torch.Tensor, # [num_blocks, num_heads, block_size, head_size]
|
|
v_cache: torch.Tensor,
|
|
context_lengths: torch.Tensor, # [num_seqs]
|
|
block_tables: torch.Tensor, # [num_seqs,max_blocks_per_sequence]
|
|
attn_mask: torch.Tensor = None, # [bsz, input_lengths + output_lengths]
|
|
):
|
|
# Firt, do shape verification
|
|
bsz, seq_len, num_heads, head_size = q.shape
|
|
num_kv_heads = k.shape[-2]
|
|
assert num_heads % num_kv_heads == 0, "num_kv_heads should be divisible by num_heads"
|
|
num_kv_groups = num_heads // num_kv_heads
|
|
block_size = k_cache.size(-2)
|
|
assert q.shape[0] == k.shape[0] == v.shape[0] == block_tables.shape[0]
|
|
block_tables.shape[-1] * block_size
|
|
|
|
# Copy kv to memory(rotary embedded)
|
|
copy_to_cache(k, k_cache, lengths=context_lengths, block_tables=block_tables)
|
|
copy_to_cache(v, v_cache, lengths=context_lengths, block_tables=block_tables)
|
|
|
|
q = q.transpose(1, 2)
|
|
k = PagedAttention.repeat_kv(k.transpose(1, 2), num_kv_groups)
|
|
v = PagedAttention.repeat_kv(v.transpose(1, 2), num_kv_groups)
|
|
|
|
attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(head_size)
|
|
|
|
padding_mask = None
|
|
|
|
if attn_mask is not None:
|
|
padding_mask = AttentionMaskConverter._expand_mask(attn_mask, q.dtype, seq_len)
|
|
|
|
attn_mask = AttentionMaskConverter._make_causal_mask(
|
|
(bsz, seq_len), q.dtype, q.device, past_key_values_length=seq_len - seq_len
|
|
)
|
|
|
|
if padding_mask is not None:
|
|
attn_mask = attn_mask.masked_fill(padding_mask.bool(), torch.finfo(q.dtype).min)
|
|
|
|
if attn_weights.size() != (bsz, num_heads, seq_len, seq_len):
|
|
raise ValueError(f"Got wrong attn_weights, should be in shape {(bsz,num_heads,seq_len,seq_len)}.")
|
|
if attn_mask is not None:
|
|
attn_weights += attn_mask
|
|
|
|
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
|
|
attn_output = torch.matmul(attn_weights, v)
|
|
|
|
if attn_output.size() != (bsz, num_heads, seq_len, head_size):
|
|
raise ValueError(f"Got wrong attn_output, should be in shape {(bsz,num_heads,seq_len,head_size)}.")
|
|
|
|
attn_output = attn_output.transpose(1, 2).contiguous().reshape(bsz, seq_len, -1)
|
|
|
|
return attn_output
|
|
|
|
@staticmethod
|
|
def pad_decoding_forward(
|
|
q: torch.Tensor, # [bsz, 1, num_heads, head_size]
|
|
k: torch.Tensor, # [bsz, 1, num_kv_heads, head_size]
|
|
v: torch.Tensor,
|
|
k_cache: torch.Tensor, # [num_blocks, num_heads, block_size, head_size]
|
|
v_cache: torch.Tensor,
|
|
lengths: torch.Tensor, # [num_seqs]: input_lengths + output_lengths
|
|
block_tables: torch.Tensor, # [num_seqs,max_blocks_per_sequence]
|
|
attn_mask: torch.Tensor = None, # [bsz, input_lengths + output_lengths]
|
|
):
|
|
# Firt, do shape verification.
|
|
bsz, q_length, num_heads, head_size = q.shape
|
|
|
|
num_kv_heads = k.shape[-2]
|
|
assert num_heads % num_kv_heads == 0, "num_kv_heads should be divisible by num_heads"
|
|
num_kv_groups = num_heads // num_kv_heads
|
|
seq_len = max(lengths)
|
|
|
|
assert q.shape[0] == k.shape[0] == v.shape[0] == block_tables.shape[0]
|
|
|
|
copy_to_cache(k, k_cache, lengths=lengths, block_tables=block_tables, type="decoding")
|
|
copy_to_cache(v, v_cache, lengths=lengths, block_tables=block_tables, type="decoding")
|
|
|
|
k = convert_kvcache(k_cache, lengths, block_tables) # bsz, seqlen,
|
|
v = convert_kvcache(v_cache, lengths, block_tables)
|
|
|
|
q = q.transpose(1, 2)
|
|
k = PagedAttention.repeat_kv(k.transpose(1, 2), num_kv_groups)
|
|
v = PagedAttention.repeat_kv(v.transpose(1, 2), num_kv_groups)
|
|
|
|
attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(head_size)
|
|
if attn_weights.size() != (bsz, num_heads, 1, seq_len):
|
|
raise ValueError(f"Got wrong attn_weights, should be in shape {(bsz,num_heads,1,seq_len)}.")
|
|
|
|
padding_mask = None
|
|
if attn_mask is not None:
|
|
padding_mask = AttentionMaskConverter._expand_mask(attn_mask, q.dtype, q_length)
|
|
|
|
attn_mask = AttentionMaskConverter._make_causal_mask(
|
|
(bsz, q_length), q.dtype, q.device, past_key_values_length=seq_len - q_length
|
|
)
|
|
|
|
if padding_mask is not None:
|
|
attn_mask = attn_mask.masked_fill(padding_mask.bool(), torch.finfo(q.dtype).min)
|
|
|
|
attn_weights += attn_mask
|
|
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
|
|
attn_output = torch.matmul(attn_weights, v)
|
|
|
|
if attn_output.size() != (bsz, num_heads, 1, head_size):
|
|
raise ValueError(f"Got wrong attn_output, should be in shape {(bsz,num_heads,1,head_size)}.")
|
|
attn_output = attn_output.transpose(1, 2).contiguous().reshape(bsz, 1, -1)
|
|
|
|
return attn_output
|
|
|
|
@staticmethod
|
|
def no_pad_decoding_forward(
|
|
self,
|
|
q: torch.Tensor, # [num_tokens, num_heads, head_size]
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
k_cache: torch.Tensor, # [num_blocks, num_heads, head_size, block_size]
|
|
v_cache: torch.Tensor,
|
|
lengths: torch.Tensor, # [num_seqs]: input_lengths + output_lengths
|
|
block_tables: torch.Tensor, # [num_seqs,max_blocks_per_sequence]
|
|
):
|
|
return self.pad_decoding_forward(
|
|
q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1), k_cache, v_cache, lengths, block_tables
|
|
)
|