mirror of
https://github.com/PaddlePaddle/FastDeploy.git
synced 2025-10-05 08:37:06 +08:00
[Iluvatar GPU] Optimze attention and moe performance (#3234)
This commit is contained in:
yzwu
@@ -85,45 +85,120 @@ class IluvatarAttnBackend(AttentionBackend):
|
||||
Which is used only for testing purpose.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
llm_config: FDConfig,
|
||||
kv_num_heads: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
):
|
||||
def __init__(self, fd_config: FDConfig, kv_num_heads: int, num_heads: int, head_dim: int):
|
||||
super().__init__()
|
||||
self.attention_metadata = IluvatarAttentionMetadata()
|
||||
self.attention_metadata.block_size = llm_config.cache_config.block_size
|
||||
assert llm_config.cache_config.enc_dec_block_num == 0, "Iluvatar does not support yet"
|
||||
self.attention_metadata.block_size = fd_config.parallel_config.block_size
|
||||
assert (
|
||||
fd_config.parallel_config.enc_dec_block_num == 0
|
||||
), f"Iluvatar does not support yet, {fd_config.parallel_config.enc_dec_block_num}"
|
||||
assert self.attention_metadata.block_size == 16, "Iluvatar paged attn requires block_size must be 16."
|
||||
|
||||
self.attention_metadata.max_context_len = llm_config.parallel_config.max_model_len
|
||||
self.attention_metadata.causal = getattr(llm_config.model_config, "causal", True)
|
||||
self.speculate_method = getattr(llm_config.parallel_config, "speculate_method", None)
|
||||
self.attention_metadata.max_context_len = fd_config.parallel_config.max_model_len
|
||||
self.attention_metadata.causal = getattr(fd_config.model_config, "causal", True)
|
||||
self.speculate_method = getattr(fd_config.parallel_config, "speculate_method", None)
|
||||
self.use_speculate = self.speculate_method is not None
|
||||
self.attention_metadata.num_kv_heads = kv_num_heads
|
||||
self.attention_metadata.dropout = llm_config.model_config.hidden_dropout_prob
|
||||
self.attention_metadata.dropout = fd_config.model_config.hidden_dropout_prob
|
||||
self.num_heads = num_heads
|
||||
self.total_num_heads = num_heads + 2 * kv_num_heads
|
||||
self.head_dim = head_dim
|
||||
self.hidden_dim = num_heads * head_dim
|
||||
self.total_hidden_dim = self.total_num_heads * head_dim
|
||||
# note: scale need to change if using MLA
|
||||
self.attention_metadata.scale = 1.0 / sqrt(head_dim)
|
||||
self.num_layers = llm_config.model_config.num_hidden_layers
|
||||
self.num_layers = fd_config.model_config.num_hidden_layers
|
||||
self.dtype = paddle.get_default_dtype()
|
||||
|
||||
self.record_block_table_metadata = {}
|
||||
self.only_use_flash_attn = int(os.getenv("FD_ILUVATAR_ONLY_USE_FLASH_ATTN", 0)) == 1
|
||||
self.do_check_kv_cache = int(os.getenv("FD_ILUVATAR_CHECK_KV_CACHE_CORRECTNESS", 0)) == 1
|
||||
if not self.only_use_flash_attn:
|
||||
assert self.attention_metadata.block_size == 16, "Iluvatar paged attn requires block_size must be 16."
|
||||
if self.do_check_kv_cache:
|
||||
self.record_batched_k = [{} for _ in range(self.num_layers)]
|
||||
self.record_batched_v = [{} for _ in range(self.num_layers)]
|
||||
self.enable_fused_attention = int(os.getenv("FD_ILUVATAR_ENABLE_FUSED_ATTN", 1))
|
||||
|
||||
def init_attention_metadata(self, forward_meta: ForwardMeta):
|
||||
"""Initialize attntion metadata hence all layers in the forward pass can reuse it."""
|
||||
self.attention_metadata.block_tables = forward_meta.block_tables
|
||||
self.attention_metadata.attn_mask = forward_meta.attn_mask
|
||||
self.attention_metadata.seq_lens = forward_meta.seq_lens_decoder
|
||||
self.attention_metadata.cu_seqlens_q = forward_meta.cu_seqlens_q
|
||||
self.attention_metadata.cu_seqlens_k = forward_meta.cu_seqlens_k
|
||||
self.prefill_info_dict = {}
|
||||
self.decode_info_dict = {}
|
||||
|
||||
prefill_non_zeros_ids = forward_meta.seq_lens_this_time > 1
|
||||
decode_non_zeros_ids = forward_meta.seq_lens_this_time == 1
|
||||
self.prefill_info_dict["batch_ids"] = paddle.where(prefill_non_zeros_ids)[0]
|
||||
self.decode_info_dict["batch_ids"] = paddle.where(decode_non_zeros_ids)[0]
|
||||
|
||||
self.prefill_len = len(self.prefill_info_dict["batch_ids"])
|
||||
self.decode_len = len(self.decode_info_dict["batch_ids"])
|
||||
# only prefill
|
||||
if self.decode_len == 0:
|
||||
cu_seq_ids = list(range(self.prefill_len + 1))
|
||||
self.prefill_info_dict["cu_seqlens_q"] = forward_meta.cu_seqlens_q[cu_seq_ids]
|
||||
# only decode
|
||||
elif self.prefill_len == 0:
|
||||
pass
|
||||
# both prefill and decode
|
||||
else:
|
||||
prefill_num_tokens = paddle.sum(forward_meta.seq_lens_this_time[prefill_non_zeros_ids])
|
||||
decode_num_tokens = paddle.sum(forward_meta.seq_lens_this_time[decode_non_zeros_ids])
|
||||
|
||||
self.prefill_info_dict["cu_seqlens_q"] = paddle.zeros(
|
||||
[self.prefill_len + 1], dtype=forward_meta.cu_seqlens_q.dtype
|
||||
)
|
||||
self.prefill_info_dict["cu_seqlens_q"][1:] = forward_meta.seq_lens_encoder[
|
||||
self.prefill_info_dict["batch_ids"], 0
|
||||
]
|
||||
self.prefill_info_dict["cu_seqlens_q"] = paddle.cumsum(self.prefill_info_dict["cu_seqlens_q"])
|
||||
|
||||
self.prefill_qkv = paddle.zeros([prefill_num_tokens, self.total_hidden_dim], dtype=self.dtype)
|
||||
self.decode_qkv = paddle.zeros([decode_num_tokens, self.total_hidden_dim], dtype=self.dtype)
|
||||
self.merged_output = paddle.zeros(
|
||||
[prefill_num_tokens + decode_num_tokens, self.num_heads, self.head_dim], dtype=self.dtype
|
||||
)
|
||||
|
||||
prefill_start, decode_start, start = 0, 0, 0
|
||||
non_zeros_ids = forward_meta.seq_lens_this_time != 0
|
||||
non_zeros_seq_lens = forward_meta.seq_lens_this_time[non_zeros_ids]
|
||||
end = non_zeros_seq_lens[0]
|
||||
if end > 1:
|
||||
last_stage = "prefill"
|
||||
prefill_end = end
|
||||
decode_end = 0
|
||||
else:
|
||||
last_stage = "decode"
|
||||
prefill_end = 0
|
||||
decode_end = end
|
||||
|
||||
self.prefill_info_dict["id_group"] = []
|
||||
self.prefill_info_dict["reverse_id_group"] = []
|
||||
self.decode_info_dict["id_group"] = []
|
||||
self.decode_info_dict["reverse_id_group"] = []
|
||||
self.record_stages = []
|
||||
for seq_len in non_zeros_seq_lens[1:]:
|
||||
if seq_len > 1:
|
||||
if last_stage == "decode":
|
||||
self.record_stages.append((last_stage, len(self.decode_info_dict["id_group"])))
|
||||
self.decode_info_dict["id_group"].append((decode_start, decode_end))
|
||||
self.decode_info_dict["reverse_id_group"].append((start, end))
|
||||
decode_start = decode_end
|
||||
start = end
|
||||
last_stage = "prefill"
|
||||
prefill_end += seq_len
|
||||
end += seq_len
|
||||
else:
|
||||
if last_stage == "prefill":
|
||||
self.record_stages.append((last_stage, len(self.prefill_info_dict["id_group"])))
|
||||
self.prefill_info_dict["id_group"].append((prefill_start, prefill_end))
|
||||
self.prefill_info_dict["reverse_id_group"].append((start, end))
|
||||
prefill_start = prefill_end
|
||||
start = end
|
||||
last_stage = "decode"
|
||||
decode_end += seq_len
|
||||
end += seq_len
|
||||
|
||||
if prefill_start < prefill_end:
|
||||
self.record_stages.append(("prefill", len(self.prefill_info_dict["id_group"])))
|
||||
self.prefill_info_dict["id_group"].append((prefill_start, prefill_end))
|
||||
self.prefill_info_dict["reverse_id_group"].append((start, end))
|
||||
if decode_start < decode_end:
|
||||
self.record_stages.append(("decode", len(self.decode_info_dict["id_group"])))
|
||||
self.decode_info_dict["id_group"].append((decode_start, decode_end))
|
||||
self.decode_info_dict["reverse_id_group"].append((start, end))
|
||||
|
||||
def get_attntion_meta(self):
|
||||
"""get_attntion_meta"""
|
||||
@@ -144,93 +219,15 @@ class IluvatarAttnBackend(AttentionBackend):
|
||||
self.head_dim,
|
||||
)
|
||||
|
||||
def get_new_kv(
|
||||
self,
|
||||
k,
|
||||
v,
|
||||
k_cache_id: int,
|
||||
v_cache_id: int,
|
||||
forward_meta: ForwardMeta,
|
||||
debug_paged_attn=False,
|
||||
):
|
||||
new_k = []
|
||||
new_v = []
|
||||
tensor_start = 0
|
||||
for batch_idx in range(forward_meta.block_tables.shape[0]):
|
||||
seq_len = forward_meta.seq_lens_this_time[batch_idx]
|
||||
if seq_len == 0:
|
||||
continue
|
||||
|
||||
tensor_end = tensor_start + seq_len
|
||||
slice_k = k[tensor_start:tensor_end, :, :]
|
||||
slice_v = v[tensor_start:tensor_end, :, :]
|
||||
|
||||
if seq_len > 1:
|
||||
# prefill
|
||||
new_k.append(slice_k)
|
||||
new_v.append(slice_v)
|
||||
else:
|
||||
# decode
|
||||
assert seq_len == 1
|
||||
cur_block_tables = forward_meta.block_tables[batch_idx]
|
||||
cur_used_block_tables = cur_block_tables[cur_block_tables != -1]
|
||||
assert (
|
||||
batch_idx in self.record_block_table_metadata
|
||||
), f"Key error: {batch_idx} vs {self.record_block_table_metadata}."
|
||||
cur_block_table_metadata = self.record_block_table_metadata[batch_idx]
|
||||
record_last_block_id = cur_block_table_metadata["block_id"]
|
||||
assert record_last_block_id != -1
|
||||
for block_id in cur_used_block_tables:
|
||||
if block_id == record_last_block_id:
|
||||
cache_end = cur_block_table_metadata["cache_end"]
|
||||
block_k_cache = forward_meta.caches[k_cache_id][block_id, :, 0:cache_end, :]
|
||||
block_v_cache = forward_meta.caches[v_cache_id][block_id, :, 0:cache_end, :]
|
||||
else:
|
||||
block_k_cache = forward_meta.caches[k_cache_id][block_id]
|
||||
block_v_cache = forward_meta.caches[v_cache_id][block_id]
|
||||
|
||||
# [num_kv_heads, block_size, head_dim] -> [block_size, num_kv_heads, head_dim]
|
||||
new_k.append(block_k_cache.transpose([1, 0, 2]).contiguous())
|
||||
new_v.append(block_v_cache.transpose([1, 0, 2]).contiguous())
|
||||
if block_id == record_last_block_id:
|
||||
break
|
||||
|
||||
# as line 301 show, record_block_table_metadata updates when executing the last layer,
|
||||
# so slice_k and slice_v has been updated in block_k_cache and block_v_cache
|
||||
if not (debug_paged_attn and (k_cache_id / 2 == self.num_layers - 1)):
|
||||
new_k.append(slice_k)
|
||||
new_v.append(slice_v)
|
||||
|
||||
tensor_start = tensor_end
|
||||
|
||||
if len(new_k) == 1:
|
||||
return new_k[0], new_v[0]
|
||||
else:
|
||||
new_k = paddle.concat(new_k, axis=0)
|
||||
new_v = paddle.concat(new_v, axis=0)
|
||||
return new_k, new_v
|
||||
|
||||
def update_kv_cache(
|
||||
self,
|
||||
k,
|
||||
v,
|
||||
k_cache_id: int,
|
||||
v_cache_id: int,
|
||||
layer_id: int,
|
||||
forward_meta: ForwardMeta,
|
||||
specific_batch_ids=None,
|
||||
debug_paged_attn=False,
|
||||
def prefill_update_kv_cache(
|
||||
self, k, v, k_cache_id: int, v_cache_id: int, layer_id: int, forward_meta: ForwardMeta, prefill_batch_ids: list
|
||||
):
|
||||
# [num_tokens, num_kv_heads, head_dim] -> [num_kv_heads, num_tokens, head_dim]
|
||||
trans_k = k.transpose([1, 0, 2]).contiguous()
|
||||
trans_v = v.transpose([1, 0, 2]).contiguous()
|
||||
tensor_start = 0
|
||||
for batch_idx in range(forward_meta.block_tables.shape[0]):
|
||||
if specific_batch_ids is not None and batch_idx not in specific_batch_ids:
|
||||
continue
|
||||
for batch_idx in prefill_batch_ids:
|
||||
seq_len = forward_meta.seq_lens_this_time[batch_idx]
|
||||
if seq_len == 0:
|
||||
continue
|
||||
|
||||
tensor_end = tensor_start + seq_len
|
||||
slice_trans_k = trans_k[:, tensor_start:tensor_end, :]
|
||||
@@ -239,146 +236,67 @@ class IluvatarAttnBackend(AttentionBackend):
|
||||
cur_block_tables = forward_meta.block_tables[batch_idx]
|
||||
cur_used_block_tables = cur_block_tables[cur_block_tables != -1]
|
||||
|
||||
# prefill
|
||||
if seq_len > 1:
|
||||
cache_start = 0
|
||||
cur_used_num_blocks = cur_used_block_tables.shape[0]
|
||||
for i, block_id in enumerate(cur_used_block_tables):
|
||||
# last block: seq_len - cache_start <= block_size
|
||||
if i == cur_used_num_blocks - 1:
|
||||
cache_end = seq_len - cache_start
|
||||
assert cache_end <= self.attention_metadata.block_size
|
||||
forward_meta.caches[k_cache_id][block_id, :, 0:cache_end, :] = slice_trans_k[
|
||||
:, cache_start:seq_len, :
|
||||
]
|
||||
forward_meta.caches[v_cache_id][block_id, :, 0:cache_end, :] = slice_trans_v[
|
||||
:, cache_start:seq_len, :
|
||||
]
|
||||
if layer_id == self.num_layers - 1:
|
||||
self.record_block_table_metadata[batch_idx] = {
|
||||
"block_id": block_id.item(),
|
||||
"cache_end": cache_end,
|
||||
}
|
||||
# non last block: seq_lens_this_time > block_size
|
||||
else:
|
||||
assert seq_len > self.attention_metadata.block_size
|
||||
cache_end = cache_start + self.attention_metadata.block_size
|
||||
forward_meta.caches[k_cache_id][block_id] = slice_trans_k[:, cache_start:cache_end, :]
|
||||
forward_meta.caches[v_cache_id][block_id] = slice_trans_v[:, cache_start:cache_end, :]
|
||||
cache_start += self.attention_metadata.block_size
|
||||
else:
|
||||
# decode
|
||||
assert seq_len == 1
|
||||
cur_last_block_id = cur_used_block_tables[-1].item()
|
||||
assert cur_last_block_id != -1
|
||||
assert (
|
||||
batch_idx in self.record_block_table_metadata
|
||||
), f"Key error: {batch_idx} vs {self.record_block_table_metadata}."
|
||||
cur_block_table_metadata = self.record_block_table_metadata[batch_idx]
|
||||
record_last_block_id = cur_block_table_metadata["block_id"]
|
||||
|
||||
if cur_last_block_id == record_last_block_id:
|
||||
# not alloc new block in decode stage
|
||||
cache_start = cur_block_table_metadata["cache_end"]
|
||||
cache_start = 0
|
||||
cur_used_num_blocks = cur_used_block_tables.shape[0]
|
||||
for i, block_id in enumerate(cur_used_block_tables):
|
||||
# last block: seq_len - cache_start <= block_size
|
||||
if i == cur_used_num_blocks - 1:
|
||||
cache_end = seq_len - cache_start
|
||||
assert cache_end <= self.attention_metadata.block_size
|
||||
paddle.assign(
|
||||
slice_trans_k[:, cache_start:seq_len, :],
|
||||
output=forward_meta.caches[k_cache_id][block_id, :, 0:cache_end, :],
|
||||
)
|
||||
paddle.assign(
|
||||
slice_trans_v[:, cache_start:seq_len, :],
|
||||
output=forward_meta.caches[v_cache_id][block_id, :, 0:cache_end, :],
|
||||
)
|
||||
if layer_id == self.num_layers - 1:
|
||||
self.record_block_table_metadata[batch_idx] = {
|
||||
"block_id": block_id.item(),
|
||||
"cache_end": cache_end.item(),
|
||||
}
|
||||
# non last block: seq_lens_this_time > block_size
|
||||
else:
|
||||
# alloc new block in decode stage
|
||||
cache_start = 0
|
||||
|
||||
cache_end = cache_start + 1
|
||||
assert cache_end <= self.attention_metadata.block_size
|
||||
|
||||
# paged attn API will update kv cache with inplace mode
|
||||
if not debug_paged_attn:
|
||||
forward_meta.caches[k_cache_id][cur_last_block_id, :, cache_start:cache_end, :] = slice_trans_k
|
||||
forward_meta.caches[v_cache_id][cur_last_block_id, :, cache_start:cache_end, :] = slice_trans_v
|
||||
|
||||
# update record_block_table_metadata
|
||||
if layer_id == self.num_layers - 1:
|
||||
self.record_block_table_metadata[batch_idx]["block_id"] = cur_last_block_id
|
||||
self.record_block_table_metadata[batch_idx]["cache_end"] = cache_end
|
||||
assert seq_len > self.attention_metadata.block_size
|
||||
cache_end = cache_start + self.attention_metadata.block_size
|
||||
paddle.assign(
|
||||
slice_trans_k[:, cache_start:cache_end, :], output=forward_meta.caches[k_cache_id][block_id]
|
||||
)
|
||||
paddle.assign(
|
||||
slice_trans_v[:, cache_start:cache_end, :], output=forward_meta.caches[v_cache_id][block_id]
|
||||
)
|
||||
cache_start += self.attention_metadata.block_size
|
||||
|
||||
tensor_start = tensor_end
|
||||
|
||||
def _check_new_kv_correctness(self, k, v, new_k, new_v, layer_id: int, forward_meta: ForwardMeta):
|
||||
tensor_start = 0
|
||||
for batch_idx, seq_lens_this_time in enumerate(forward_meta.seq_lens_this_time):
|
||||
if seq_lens_this_time == 0:
|
||||
continue
|
||||
# note: the second request will also use the batch_idx 0 instead of 1 in
|
||||
# the streaming inference mode, so use seq_lens_this_time > 1 with the same
|
||||
# batch_idx represents the second request comes.
|
||||
if seq_lens_this_time > 1 and batch_idx in self.record_batched_k[layer_id]:
|
||||
print(
|
||||
f"clear self.record_batched_batched_k: "
|
||||
f"layer_id={layer_id}, batch_id={batch_idx}, "
|
||||
f"record_lens={len(self.record_batched_k[layer_id][batch_idx])}"
|
||||
)
|
||||
self.record_batched_k[layer_id][batch_idx].clear()
|
||||
self.record_batched_v[layer_id][batch_idx].clear()
|
||||
tensor_end = tensor_start + seq_lens_this_time
|
||||
slice_k = k[tensor_start:tensor_end, :, :]
|
||||
slice_v = v[tensor_start:tensor_end, :, :]
|
||||
if batch_idx not in self.record_batched_k[layer_id]:
|
||||
self.record_batched_k[layer_id][batch_idx] = []
|
||||
self.record_batched_v[layer_id][batch_idx] = []
|
||||
self.record_batched_k[layer_id][batch_idx].append(slice_k)
|
||||
self.record_batched_v[layer_id][batch_idx].append(slice_v)
|
||||
tensor_start = tensor_end
|
||||
|
||||
ref_k, ref_v = [], []
|
||||
for batch_idx, seq_lens_this_time in enumerate(forward_meta.seq_lens_this_time):
|
||||
if seq_lens_this_time == 0:
|
||||
continue
|
||||
bached_k_list = self.record_batched_k[layer_id][batch_idx]
|
||||
bached_v_list = self.record_batched_v[layer_id][batch_idx]
|
||||
ref_k.extend(bached_k_list)
|
||||
ref_v.extend(bached_v_list)
|
||||
|
||||
ref_k = paddle.concat(ref_k, axis=0)
|
||||
ref_v = paddle.concat(ref_v, axis=0)
|
||||
print(
|
||||
f"_check_new_kv_correctness: layer_id={layer_id}, "
|
||||
f"k.shape={k.shape}, v.shape={v.shape}, "
|
||||
f"ref_k.shape={ref_k.shape}, ref_v.shape={ref_v.shape}, "
|
||||
f"new_k.shape={new_k.shape}, new_v.shape={new_v.shape}, "
|
||||
f"len(self.record_batched_k[layer_id])={len(self.record_batched_k[layer_id])}, "
|
||||
f"len(self.record_batched_k[layer_id][0])={len(self.record_batched_k[layer_id][0])}, "
|
||||
f"forward_meta.seq_lens_this_time={forward_meta.seq_lens_this_time}"
|
||||
f"ref_k[-2:, 0:2, 0:2]={ref_k[-2:, 0:2, 0:2]}, "
|
||||
f"ref_v[-2:, 0:2, 0:2]={ref_v[-2:, 0:2, 0:2]}, "
|
||||
f"new_k[-2:, 0:2, 0:2]={new_k[-2:, 0:2, 0:2]}, "
|
||||
f"new_v[-2:, 0:2, 0:2]={new_v[-2:, 0:2, 0:2]}"
|
||||
)
|
||||
assert paddle.allclose(
|
||||
ref_k.to("cpu").to(paddle.float32),
|
||||
new_k.to("cpu").to(paddle.float32),
|
||||
)
|
||||
assert paddle.allclose(
|
||||
ref_v.to("cpu").to(paddle.float32),
|
||||
new_v.to("cpu").to(paddle.float32),
|
||||
)
|
||||
|
||||
def get_splited_qkv(self, qkv: paddle.Tensor, forward_meta: ForwardMeta):
|
||||
q_end = self.num_heads * self.head_dim
|
||||
def get_splited_qkv(
|
||||
self, qkv: paddle.Tensor, forward_meta: ForwardMeta, cu_seqlens_q: paddle.Tensor, batch_ids=None
|
||||
):
|
||||
q_end = self.hidden_dim
|
||||
k_end = q_end + self.attention_metadata.num_kv_heads * self.head_dim
|
||||
v_end = k_end + self.attention_metadata.num_kv_heads * self.head_dim
|
||||
assert v_end == qkv.shape[-1], f"Shape mistach: {v_end} vs {qkv.shape[-1]}"
|
||||
assert qkv.shape[0] == forward_meta.cu_seqlens_q[-1]
|
||||
assert v_end == qkv.shape[-1], f"Shape mismatch: {v_end} vs {qkv.shape[-1]}"
|
||||
assert qkv.shape[0] == cu_seqlens_q[-1], f"Shape mismatch: {qkv.shape[0]} vs {cu_seqlens_q[-1]}"
|
||||
|
||||
if batch_ids is None:
|
||||
batch_ids = list(range(forward_meta.seq_lens_this_time.shape[0]))
|
||||
|
||||
q = qkv[..., 0:q_end]
|
||||
k = qkv[..., q_end:k_end]
|
||||
v = qkv[..., k_end:v_end]
|
||||
q = q.view([-1, self.num_heads, self.head_dim]).contiguous()
|
||||
k = k.view([-1, self.attention_metadata.num_kv_heads, self.head_dim]).contiguous()
|
||||
v = v.view([-1, self.attention_metadata.num_kv_heads, self.head_dim]).contiguous()
|
||||
# forward_meta.seq_lens_this_time [max_batch,]
|
||||
for batch_idx in range(forward_meta.seq_lens_this_time.shape[0]):
|
||||
q = q.view([-1, self.num_heads, self.head_dim])
|
||||
k = k.view([-1, self.attention_metadata.num_kv_heads, self.head_dim])
|
||||
v = v.view([-1, self.attention_metadata.num_kv_heads, self.head_dim])
|
||||
|
||||
for idx in range(len(cu_seqlens_q) - 1):
|
||||
batch_idx = batch_ids[idx]
|
||||
seq_len_i = forward_meta.seq_lens_this_time[batch_idx]
|
||||
if seq_len_i == 0:
|
||||
continue
|
||||
cached_kv_len = forward_meta.seq_lens_decoder[batch_idx][0]
|
||||
cu_seq_start_q = forward_meta.cu_seqlens_q[batch_idx]
|
||||
cu_seq_end_q = forward_meta.cu_seqlens_q[batch_idx + 1]
|
||||
cu_seq_start_q = cu_seqlens_q[idx]
|
||||
cu_seq_end_q = cu_seqlens_q[idx + 1]
|
||||
# forward_meta.rotary_embs is [2, 1, S, 1, D]
|
||||
if forward_meta.rotary_embs is not None:
|
||||
cos = forward_meta.rotary_embs[0, 0, cached_kv_len : cached_kv_len + seq_len_i, :, :]
|
||||
@@ -388,75 +306,114 @@ class IluvatarAttnBackend(AttentionBackend):
|
||||
|
||||
return q, k, v
|
||||
|
||||
def get_splited_info_by_stage(self, q, k, v, forward_meta: ForwardMeta):
|
||||
prefill_info_dict = {"q": [], "k": [], "v": [], "batch_ids": []}
|
||||
decode_info_dict = {"q": [], "k": [], "v": [], "batch_ids": []}
|
||||
tensor_start = 0
|
||||
for batch_idx, seq_lens_this_time in enumerate(forward_meta.seq_lens_this_time):
|
||||
if seq_lens_this_time == 0:
|
||||
continue
|
||||
tensor_end = tensor_start + seq_lens_this_time
|
||||
slice_q = q[tensor_start:tensor_end, :, :]
|
||||
slice_k = k[tensor_start:tensor_end, :, :]
|
||||
slice_v = v[tensor_start:tensor_end, :, :]
|
||||
if seq_lens_this_time > 1:
|
||||
prefill_info_dict["q"].append(slice_q)
|
||||
prefill_info_dict["k"].append(slice_k)
|
||||
prefill_info_dict["v"].append(slice_v)
|
||||
prefill_info_dict["batch_ids"].append(batch_idx)
|
||||
def split_pd_qkv(self, qkv):
|
||||
|
||||
for ids, reverse_ids in zip(self.prefill_info_dict["id_group"], self.prefill_info_dict["reverse_id_group"]):
|
||||
self.prefill_qkv[ids[0] : ids[1], :] = qkv[reverse_ids[0] : reverse_ids[1], :]
|
||||
|
||||
for ids, reverse_ids in zip(self.decode_info_dict["id_group"], self.decode_info_dict["reverse_id_group"]):
|
||||
self.decode_qkv[ids[0] : ids[1], :] = qkv[reverse_ids[0] : reverse_ids[1], :]
|
||||
|
||||
return self.prefill_qkv, self.decode_qkv
|
||||
|
||||
def merge_pd_output(self, prefill_out, decode_out):
|
||||
for stage, idx in self.record_stages:
|
||||
if stage == "prefill":
|
||||
ids = self.prefill_info_dict["id_group"][idx]
|
||||
reverse_ids = self.prefill_info_dict["reverse_id_group"][idx]
|
||||
self.merged_output[reverse_ids[0] : reverse_ids[1], :, :] = prefill_out[ids[0] : ids[1], :, :]
|
||||
else:
|
||||
assert seq_lens_this_time == 1
|
||||
decode_info_dict["q"].append(slice_q)
|
||||
decode_info_dict["k"].append(slice_k)
|
||||
decode_info_dict["v"].append(slice_v)
|
||||
decode_info_dict["batch_ids"].append(batch_idx)
|
||||
tensor_start = tensor_end
|
||||
ids = self.decode_info_dict["id_group"][idx]
|
||||
reverse_ids = self.decode_info_dict["reverse_id_group"][idx]
|
||||
self.merged_output[reverse_ids[0] : reverse_ids[1], :, :] = decode_out[ids[0] : ids[1], :, :]
|
||||
return self.merged_output
|
||||
|
||||
if len(prefill_info_dict["batch_ids"]) > 0:
|
||||
prefill_info_dict["q"] = paddle.concat(prefill_info_dict["q"], axis=0)
|
||||
prefill_info_dict["k"] = paddle.concat(prefill_info_dict["k"], axis=0)
|
||||
prefill_info_dict["v"] = paddle.concat(prefill_info_dict["v"], axis=0)
|
||||
cu_seq_ids = list(map(lambda x: x + 1, prefill_info_dict["batch_ids"]))
|
||||
prefill_info_dict["cu_seq_ids"] = [0, *cu_seq_ids]
|
||||
def forward_prefill(self, prefill_qkv, layer_id, k_cache_id, v_cache_id, forward_meta: ForwardMeta):
|
||||
prefill_q, prefill_k, prefill_v = self.get_splited_qkv(
|
||||
prefill_qkv,
|
||||
forward_meta,
|
||||
self.prefill_info_dict["cu_seqlens_q"],
|
||||
batch_ids=self.prefill_info_dict["batch_ids"],
|
||||
)
|
||||
|
||||
if len(decode_info_dict["batch_ids"]) > 0:
|
||||
decode_info_dict["q"] = paddle.concat(decode_info_dict["q"], axis=0)
|
||||
decode_info_dict["k"] = paddle.concat(decode_info_dict["k"], axis=0)
|
||||
decode_info_dict["v"] = paddle.concat(decode_info_dict["v"], axis=0)
|
||||
prefill_out = flash_attn_unpadded(
|
||||
prefill_q,
|
||||
prefill_k,
|
||||
prefill_v,
|
||||
cu_seqlens_q=self.prefill_info_dict["cu_seqlens_q"],
|
||||
cu_seqlens_k=self.prefill_info_dict["cu_seqlens_q"],
|
||||
max_seqlen_q=self.attention_metadata.max_context_len,
|
||||
max_seqlen_k=self.attention_metadata.max_context_len,
|
||||
scale=self.attention_metadata.scale,
|
||||
dropout=self.attention_metadata.dropout,
|
||||
causal=self.attention_metadata.causal,
|
||||
return_softmax=self.attention_metadata.return_softmax,
|
||||
)[0]
|
||||
self.prefill_update_kv_cache(
|
||||
prefill_k, prefill_v, k_cache_id, v_cache_id, layer_id, forward_meta, self.prefill_info_dict["batch_ids"]
|
||||
)
|
||||
|
||||
return prefill_info_dict, decode_info_dict
|
||||
return prefill_out
|
||||
|
||||
def merge_output(self, prefill_out, decode_out, forward_meta: ForwardMeta):
|
||||
assert not (prefill_out is None and decode_out is None), "prefill and decode output cannot both be None"
|
||||
if prefill_out is None:
|
||||
return decode_out
|
||||
elif decode_out is None:
|
||||
return prefill_out
|
||||
def forward_decode(self, decode_qkv, k_cache_id, v_cache_id, forward_meta: ForwardMeta):
|
||||
k_cache = forward_meta.caches[k_cache_id]
|
||||
v_cache = forward_meta.caches[v_cache_id]
|
||||
if self.enable_fused_attention:
|
||||
rope_cos = forward_meta.rotary_embs[0, 0, :, :, :]
|
||||
rope_sin = forward_meta.rotary_embs[1, 0, :, :, :]
|
||||
decode_out = paged_attention(
|
||||
decode_qkv.view([-1, self.total_num_heads, self.head_dim]),
|
||||
k_cache,
|
||||
v_cache,
|
||||
block_tables=forward_meta.block_tables[self.decode_info_dict["batch_ids"], :],
|
||||
seq_lens=forward_meta.seq_lens_decoder[self.decode_info_dict["batch_ids"], 0] + 1,
|
||||
num_kv_heads=self.attention_metadata.num_kv_heads,
|
||||
scale=self.attention_metadata.scale,
|
||||
block_size=self.attention_metadata.block_size,
|
||||
max_context_len=self.attention_metadata.max_context_len,
|
||||
alibi_slopes=self.attention_metadata.alibi_slopes,
|
||||
causal=self.attention_metadata.causal,
|
||||
window_left=self.attention_metadata.window_left,
|
||||
window_right=self.attention_metadata.window_right,
|
||||
softcap=self.attention_metadata.softcap,
|
||||
use_cuda_graph=self.attention_metadata.use_cuda_graph,
|
||||
use_sqrt_alibi=self.attention_metadata.use_sqrt_alibi,
|
||||
merged_qkv=True,
|
||||
k=decode_qkv,
|
||||
v=decode_qkv,
|
||||
rope_sin=rope_sin,
|
||||
rope_cos=rope_cos,
|
||||
)
|
||||
else:
|
||||
merged_output = []
|
||||
prefill_tensor_start = 0
|
||||
decode_tensor_start = 0
|
||||
for seq_lens_this_time in forward_meta.seq_lens_this_time:
|
||||
if seq_lens_this_time == 0:
|
||||
continue
|
||||
if seq_lens_this_time > 1:
|
||||
tensor_end = prefill_tensor_start + seq_lens_this_time
|
||||
merged_output.append(prefill_out[prefill_tensor_start:tensor_end, :, :])
|
||||
prefill_tensor_start = tensor_end
|
||||
else:
|
||||
assert seq_lens_this_time == 1
|
||||
tensor_end = decode_tensor_start + seq_lens_this_time
|
||||
merged_output.append(decode_out[decode_tensor_start:tensor_end, :, :])
|
||||
decode_tensor_start = tensor_end
|
||||
decode_q, decode_k, decode_v = self.get_splited_qkv(
|
||||
decode_qkv,
|
||||
forward_meta,
|
||||
self.decode_info_dict["cu_seqlens_q"],
|
||||
batch_ids=self.decode_info_dict["batch_ids"],
|
||||
)
|
||||
|
||||
assert (
|
||||
prefill_tensor_start == prefill_out.shape[0]
|
||||
), f"prefill merged unfinished: {prefill_tensor_start} vs {prefill_out.shape[0]}"
|
||||
assert (
|
||||
decode_tensor_start == decode_out.shape[0]
|
||||
), f"decode merged unfinished: {decode_tensor_start} vs {decode_out.shape[0]}"
|
||||
merged_output = paddle.concat(merged_output, axis=0)
|
||||
return merged_output
|
||||
decode_out = paged_attention(
|
||||
decode_q,
|
||||
k_cache,
|
||||
v_cache,
|
||||
block_tables=forward_meta.block_tables[self.decode_info_dict["batch_ids"], :],
|
||||
seq_lens=forward_meta.seq_lens_decoder[self.decode_info_dict["batch_ids"], 0] + 1,
|
||||
num_kv_heads=self.attention_metadata.num_kv_heads,
|
||||
scale=self.attention_metadata.scale,
|
||||
block_size=self.attention_metadata.block_size,
|
||||
max_context_len=self.attention_metadata.max_context_len,
|
||||
alibi_slopes=self.attention_metadata.alibi_slopes,
|
||||
causal=self.attention_metadata.causal,
|
||||
window_left=self.attention_metadata.window_left,
|
||||
window_right=self.attention_metadata.window_right,
|
||||
softcap=self.attention_metadata.softcap,
|
||||
use_cuda_graph=self.attention_metadata.use_cuda_graph,
|
||||
use_sqrt_alibi=self.attention_metadata.use_sqrt_alibi,
|
||||
k=decode_k,
|
||||
v=decode_v,
|
||||
)
|
||||
|
||||
return decode_out
|
||||
|
||||
def forward_mixed(
|
||||
self,
|
||||
@@ -476,110 +433,19 @@ class IluvatarAttnBackend(AttentionBackend):
|
||||
layer_id = layer.layer_id
|
||||
k_cache_id = layer_id * 2
|
||||
v_cache_id = k_cache_id + 1
|
||||
|
||||
assert qkv is not None
|
||||
q_dim = qkv.dim()
|
||||
q, k, v = self.get_splited_qkv(qkv, forward_meta)
|
||||
assert q_dim == 2
|
||||
|
||||
if self.only_use_flash_attn:
|
||||
new_k, new_v = self.get_new_kv(k, v, k_cache_id, v_cache_id, forward_meta)
|
||||
if self.do_check_kv_cache:
|
||||
self._check_new_kv_correctness(k, v, new_k, new_v, layer_id, forward_meta)
|
||||
if self.decode_len == 0:
|
||||
output = self.forward_prefill(qkv, layer_id, k_cache_id, v_cache_id, forward_meta)
|
||||
|
||||
out = flash_attn_unpadded(
|
||||
q,
|
||||
new_k,
|
||||
new_v,
|
||||
cu_seqlens_q=self.attention_metadata.cu_seqlens_q,
|
||||
cu_seqlens_k=self.attention_metadata.cu_seqlens_k,
|
||||
max_seqlen_q=self.attention_metadata.max_context_len,
|
||||
max_seqlen_k=self.attention_metadata.max_context_len,
|
||||
scale=self.attention_metadata.scale,
|
||||
dropout=self.attention_metadata.dropout,
|
||||
causal=self.attention_metadata.causal,
|
||||
return_softmax=self.attention_metadata.return_softmax,
|
||||
)[0]
|
||||
|
||||
self.update_kv_cache(k, v, k_cache_id, v_cache_id, layer_id, forward_meta)
|
||||
elif self.prefill_len == 0:
|
||||
output = self.forward_decode(qkv, k_cache_id, v_cache_id, forward_meta)
|
||||
else:
|
||||
prefill_info_dict, decode_info_dict = self.get_splited_info_by_stage(q, k, v, forward_meta)
|
||||
prefill_out, decode_out = None, None
|
||||
prefill_qkv, decode_qkv = self.split_pd_qkv(qkv)
|
||||
prefill_output = self.forward_prefill(prefill_qkv, layer_id, k_cache_id, v_cache_id, forward_meta)
|
||||
decode_output = self.forward_decode(decode_qkv, k_cache_id, v_cache_id, forward_meta)
|
||||
output = self.merge_pd_output(prefill_output, decode_output)
|
||||
|
||||
if len(prefill_info_dict["batch_ids"]) > 0:
|
||||
prefill_out = flash_attn_unpadded(
|
||||
prefill_info_dict["q"],
|
||||
prefill_info_dict["k"],
|
||||
prefill_info_dict["v"],
|
||||
cu_seqlens_q=forward_meta.cu_seqlens_q[prefill_info_dict["cu_seq_ids"]],
|
||||
cu_seqlens_k=forward_meta.cu_seqlens_k[prefill_info_dict["cu_seq_ids"]],
|
||||
max_seqlen_q=self.attention_metadata.max_context_len,
|
||||
max_seqlen_k=self.attention_metadata.max_context_len,
|
||||
scale=self.attention_metadata.scale,
|
||||
dropout=self.attention_metadata.dropout,
|
||||
causal=self.attention_metadata.causal,
|
||||
return_softmax=self.attention_metadata.return_softmax,
|
||||
)[0]
|
||||
self.update_kv_cache(
|
||||
prefill_info_dict["k"],
|
||||
prefill_info_dict["v"],
|
||||
k_cache_id,
|
||||
v_cache_id,
|
||||
layer_id,
|
||||
forward_meta,
|
||||
specific_batch_ids=prefill_info_dict["batch_ids"],
|
||||
)
|
||||
|
||||
if len(decode_info_dict["batch_ids"]) > 0:
|
||||
k_cache = forward_meta.caches[k_cache_id]
|
||||
v_cache = forward_meta.caches[v_cache_id]
|
||||
|
||||
decode_out = paged_attention(
|
||||
decode_info_dict["q"],
|
||||
k_cache,
|
||||
v_cache,
|
||||
block_tables=forward_meta.block_tables[decode_info_dict["batch_ids"], :],
|
||||
seq_lens=forward_meta.seq_lens_decoder[decode_info_dict["batch_ids"], 0] + 1,
|
||||
num_kv_heads=self.attention_metadata.num_kv_heads,
|
||||
scale=self.attention_metadata.scale,
|
||||
block_size=self.attention_metadata.block_size,
|
||||
max_context_len=self.attention_metadata.max_context_len,
|
||||
alibi_slopes=self.attention_metadata.alibi_slopes,
|
||||
causal=self.attention_metadata.causal,
|
||||
window_left=self.attention_metadata.window_left,
|
||||
window_right=self.attention_metadata.window_right,
|
||||
softcap=self.attention_metadata.softcap,
|
||||
use_cuda_graph=self.attention_metadata.use_cuda_graph,
|
||||
use_sqrt_alibi=self.attention_metadata.use_sqrt_alibi,
|
||||
k=decode_info_dict["k"],
|
||||
v=decode_info_dict["v"],
|
||||
)
|
||||
|
||||
if self.do_check_kv_cache:
|
||||
self.update_kv_cache(
|
||||
decode_info_dict["k"],
|
||||
decode_info_dict["v"],
|
||||
k_cache_id,
|
||||
v_cache_id,
|
||||
layer_id,
|
||||
forward_meta,
|
||||
specific_batch_ids=decode_info_dict["batch_ids"],
|
||||
debug_paged_attn=True,
|
||||
)
|
||||
|
||||
if self.do_check_kv_cache:
|
||||
new_k, new_v = self.get_new_kv(
|
||||
k,
|
||||
v,
|
||||
k_cache_id,
|
||||
v_cache_id,
|
||||
forward_meta,
|
||||
debug_paged_attn=True,
|
||||
)
|
||||
self._check_new_kv_correctness(k, v, new_k, new_v, layer_id, forward_meta)
|
||||
|
||||
out = self.merge_output(prefill_out, decode_out, forward_meta)
|
||||
|
||||
if q_dim == 2:
|
||||
out = out.view([-1, self.num_heads * self.head_dim])
|
||||
|
||||
return out
|
||||
output = output.view([-1, self.num_heads * self.head_dim])
|
||||
return output
|
||||
|
||||
Reference in New Issue
Block a user