apps/FastDeploy
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/FastDeploy.git synced 2025-10-05 08:37:06 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

polish code with new pre-commit rule (#2923)

Browse Source
This commit is contained in:
Zero Rains
2025-07-19 23:19:27 +08:00
committed by GitHub
parent b8676d71a8
commit 25698d56d1
424 changed files with 14307 additions and 13518 deletions
Download Patch File Download Diff File Expand all files Collapse all files

308
fastdeploy/model_executor/layers/attention/iluvatar_attn_backend.py
View File

@@ -17,19 +17,21 @@
from __future__ import annotations
import os
import paddle
from dataclasses import dataclass
from typing import Optional, TYPE_CHECKING
from math import sqrt
from typing import TYPE_CHECKING, Optional
import paddle
from paddle.nn.functional.flash_attention import flash_attn_unpadded
from fastdeploy.model_executor.ops.iluvatar import paged_attention
from fastdeploy.config import FDConfig
from fastdeploy.model_executor.layers.attention.attention import Attention
from fastdeploy.model_executor.layers.attention.base_attention_backend import (
AttentionBackend, AttentionMetadata)
AttentionBackend,
AttentionMetadata,
)
from fastdeploy.model_executor.ops.iluvatar import paged_attention
if TYPE_CHECKING:
from fastdeploy.model_executor.forward_meta import ForwardMeta
@@ -39,6 +41,7 @@ class IluvatarAttentionMetadata(AttentionMetadata):
"""
IluvatarAttentionMetadata
"""
# flash_attn metadata
cu_seqlens_q: Optional[paddle.Tensor] = None
cu_seqlens_k: Optional[paddle.Tensor] = None
@@ -72,8 +75,7 @@ def apply_rope(qk, cos, sin):
paddle.stack([-qk[..., 1::2], qk[..., 0::2]], axis=-1),
paddle.shape(qk),
)
out = paddle.add(paddle.multiply(qk, cos),
paddle.multiply(rotate_half, sin))
out = paddle.add(paddle.multiply(qk, cos), paddle.multiply(rotate_half, sin))
return paddle.cast(out, qk.dtype)
@@ -83,18 +85,21 @@ 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,
llm_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.parallel_config.block_size
assert llm_config.parallel_config.enc_dec_block_num == 0, "Iluvatar does not support yet"
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.causal = getattr(llm_config.model_config, "causal", True)
self.speculate_method = getattr(llm_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
@@ -104,10 +109,8 @@ class IluvatarAttnBackend(AttentionBackend):
self.attention_metadata.scale = 1.0 / sqrt(head_dim)
self.num_layers = llm_config.model_config.num_hidden_layers
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
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:
@@ -133,16 +136,22 @@ class IluvatarAttnBackend(AttentionBackend):
"""
Caculate kv cache shape
"""
return (max_num_blocks, self.attention_metadata.num_kv_heads,
self.attention_metadata.block_size, self.head_dim)
return (
max_num_blocks,
self.attention_metadata.num_kv_heads,
self.attention_metadata.block_size,
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):
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
@@ -163,39 +172,31 @@ class IluvatarAttnBackend(AttentionBackend):
# 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]
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, :]
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]
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())
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)):
if not (debug_paged_attn and (k_cache_id / 2 == self.num_layers - 1)):
new_k.append(slice_k)
new_v.append(slice_v)
@@ -208,15 +209,17 @@ class IluvatarAttnBackend(AttentionBackend):
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 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,
):
# [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()
@@ -244,39 +247,33 @@ class IluvatarAttnBackend(AttentionBackend):
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, :]
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
"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, :]
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]
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:
@@ -291,34 +288,25 @@ class IluvatarAttnBackend(AttentionBackend):
# 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
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
self.record_block_table_metadata[batch_idx]["block_id"] = cur_last_block_id
self.record_block_table_metadata[batch_idx]["cache_end"] = cache_end
tensor_start = tensor_end
def _check_new_kv_correctness(self, k, v, new_k, new_v, layer_id: int,
forward_meta: ForwardMeta):
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):
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]:
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}, "
@@ -337,8 +325,7 @@ class IluvatarAttnBackend(AttentionBackend):
tensor_start = tensor_end
ref_k, ref_v = [], []
for batch_idx, seq_lens_this_time in enumerate(
forward_meta.seq_lens_this_time):
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]
@@ -359,30 +346,30 @@ class IluvatarAttnBackend(AttentionBackend):
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]}")
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))
new_k.to("cpu").to(paddle.float32),
)
assert paddle.allclose(
ref_v.to("cpu").to(paddle.float32),
new_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
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 v_end == qkv.shape[-1], f"Shape mistach: {v_end} vs {qkv.shape[-1]}"
assert qkv.shape[0] == forward_meta.cu_seqlens_q[-1]
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()
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]):
seq_len_i = forward_meta.seq_lens_this_time[batch_idx]
@@ -393,16 +380,10 @@ class IluvatarAttnBackend(AttentionBackend):
cu_seq_end_q = forward_meta.cu_seqlens_q[batch_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, :, :]
sin = forward_meta.rotary_embs[1, 0,
cached_kv_len:cached_kv_len +
seq_len_i, :, :]
q[cu_seq_start_q:cu_seq_end_q] = apply_rope(
q[cu_seq_start_q:cu_seq_end_q], cos, sin)
k[cu_seq_start_q:cu_seq_end_q] = apply_rope(
k[cu_seq_start_q:cu_seq_end_q], cos, sin)
cos = forward_meta.rotary_embs[0, 0, cached_kv_len : cached_kv_len + seq_len_i, :, :]
sin = forward_meta.rotary_embs[1, 0, cached_kv_len : cached_kv_len + seq_len_i, :, :]
q[cu_seq_start_q:cu_seq_end_q] = apply_rope(q[cu_seq_start_q:cu_seq_end_q], cos, sin)
k[cu_seq_start_q:cu_seq_end_q] = apply_rope(k[cu_seq_start_q:cu_seq_end_q], cos, sin)
return q, k, v
@@ -410,8 +391,7 @@ class IluvatarAttnBackend(AttentionBackend):
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):
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
@@ -432,29 +412,21 @@ class IluvatarAttnBackend(AttentionBackend):
tensor_start = tensor_end
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["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]
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)
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)
return prefill_info_dict, decode_info_dict
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"
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:
@@ -468,20 +440,20 @@ class IluvatarAttnBackend(AttentionBackend):
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, :, :])
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, :, :])
merged_output.append(decode_out[decode_tensor_start:tensor_end, :, :])
decode_tensor_start = tensor_end
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]}"
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
@@ -509,11 +481,9 @@ class IluvatarAttnBackend(AttentionBackend):
q, k, v = self.get_splited_qkv(qkv, forward_meta)
if self.only_use_flash_attn:
new_k, new_v = self.get_new_kv(k, v, k_cache_id, v_cache_id,
forward_meta)
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)
self._check_new_kv_correctness(k, v, new_k, new_v, layer_id, forward_meta)
out = flash_attn_unpadded(
q,
@@ -526,13 +496,12 @@ class IluvatarAttnBackend(AttentionBackend):
scale=self.attention_metadata.scale,
dropout=self.attention_metadata.dropout,
causal=self.attention_metadata.causal,
return_softmax=self.attention_metadata.return_softmax)[0]
return_softmax=self.attention_metadata.return_softmax,
)[0]
self.update_kv_cache(k, v, k_cache_id, v_cache_id, layer_id,
forward_meta)
self.update_kv_cache(k, v, k_cache_id, v_cache_id, layer_id, forward_meta)
else:
prefill_info_dict, decode_info_dict = self.get_splited_info_by_stage(
q, k, v, forward_meta)
prefill_info_dict, decode_info_dict = self.get_splited_info_by_stage(q, k, v, forward_meta)
prefill_out, decode_out = None, None
if len(prefill_info_dict["batch_ids"]) > 0:
@@ -540,16 +509,15 @@ class IluvatarAttnBackend(AttentionBackend):
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"]],
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]
return_softmax=self.attention_metadata.return_softmax,
)[0]
self.update_kv_cache(
prefill_info_dict["k"],
prefill_info_dict["v"],
@@ -557,7 +525,8 @@ class IluvatarAttnBackend(AttentionBackend):
v_cache_id,
layer_id,
forward_meta,
specific_batch_ids=prefill_info_dict['batch_ids'])
specific_batch_ids=prefill_info_dict["batch_ids"],
)
if len(decode_info_dict["batch_ids"]) > 0:
k_cache = forward_meta.caches[k_cache_id]
@@ -567,10 +536,8 @@ class IluvatarAttnBackend(AttentionBackend):
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,
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,
@@ -583,28 +550,31 @@ class IluvatarAttnBackend(AttentionBackend):
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"])
v=decode_info_dict["v"],
)
if self.do_check_kv_cache:
self.update_kv_cache(
decode_info_dict['k'],
decode_info_dict['v'],
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)
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)
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)