153 lines
5.6 KiB
Python
153 lines
5.6 KiB
Python
import torch
|
|
from vllm.config import get_current_vllm_config_or_none
|
|
|
|
from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
|
|
from vllm_ascend.worker.kvcomp_utils import (
|
|
KVCompMetaData,
|
|
recover_request_lengths,
|
|
)
|
|
|
|
|
|
def build_kvcomp_metadata(
|
|
kvcomp_meta: KVCompMetaData,
|
|
common_meta: AscendCommonAttentionMetadata,
|
|
) -> None:
|
|
num_reqs = common_meta.num_reqs
|
|
kvcomp_meta.num_actual_tokens = common_meta.num_actual_tokens
|
|
|
|
kvcomp_meta.slot_mapping = common_meta.slot_mapping
|
|
kvcomp_meta.seq_lens_gpu = common_meta.seq_lens[:num_reqs]
|
|
|
|
real_batch_size = kvcomp_meta.seq_lens_gpu.shape[0]
|
|
assert num_reqs == real_batch_size, "the len of seq_lens_gpu is not equal with batch_size"
|
|
|
|
kvcomp_meta.actual_query_lens = recover_request_lengths(common_meta.query_start_loc[: num_reqs + 1]).to(torch.int32)
|
|
|
|
runtime_seq_lens_list = kvcomp_meta.seq_lens_gpu.tolist()
|
|
|
|
if kvcomp_meta.num_actual_tokens < real_batch_size:
|
|
runtime_seq_lens_list = runtime_seq_lens_list[: kvcomp_meta.num_actual_tokens] + [0] * (
|
|
real_batch_size - kvcomp_meta.num_actual_tokens
|
|
)
|
|
kvcomp_meta.slot_mapping[kvcomp_meta.num_actual_tokens :] = -1
|
|
|
|
runtime_max_len = max(runtime_seq_lens_list) if runtime_seq_lens_list else 0
|
|
kvcomp_meta.max_seq_len_for_hamming = (
|
|
common_meta.max_seq_len if common_meta.max_seq_len is not None else runtime_max_len
|
|
)
|
|
kvcomp_meta.block_tables_for_hamming = common_meta.block_table_tensor[:real_batch_size]
|
|
|
|
top_k_cpu = kvcomp_meta.topk_for_hamming_full_cpu[:real_batch_size].clone()
|
|
if kvcomp_meta.num_actual_tokens < real_batch_size:
|
|
top_k_cpu[kvcomp_meta.num_actual_tokens :] = 0
|
|
|
|
runtime_seq_lens_cpu = torch.tensor(runtime_seq_lens_list, dtype=torch.int32)
|
|
chunk_size = kvcomp_meta.kvcomp_config.chunk_size
|
|
remainder = runtime_seq_lens_cpu % chunk_size
|
|
|
|
new_seq_lens = torch.where(
|
|
remainder == 0,
|
|
chunk_size * top_k_cpu,
|
|
chunk_size * (top_k_cpu - 1) + remainder,
|
|
)
|
|
kvcomp_meta.seq_lens_from_hamming = new_seq_lens.tolist()
|
|
|
|
q_start_loc_slice = common_meta.query_start_loc[:real_batch_size].to(torch.int64)
|
|
|
|
torch.where(
|
|
kvcomp_meta.slot_mapping[:real_batch_size] >= 0,
|
|
torch.ones_like(q_start_loc_slice, dtype=torch.bool),
|
|
torch.zeros_like(q_start_loc_slice, dtype=torch.bool),
|
|
out=kvcomp_meta.valid_query_mask[:real_batch_size],
|
|
)
|
|
torch.where(
|
|
kvcomp_meta.valid_query_mask[:real_batch_size],
|
|
kvcomp_meta.actual_query_lens[:real_batch_size],
|
|
torch.zeros_like(kvcomp_meta.actual_query_lens[:real_batch_size]),
|
|
out=kvcomp_meta.seq_lens_for_reshape[:real_batch_size],
|
|
)
|
|
common_meta.kvcomp_metadata = kvcomp_meta
|
|
|
|
|
|
def reshape_and_cache_kvcomp(kvcomp_meta: KVCompMetaData | None, layer_index: int | None, key: torch.Tensor):
|
|
assert kvcomp_meta is not None
|
|
assert layer_index is not None
|
|
|
|
if kvcomp_meta.hashk_caches[layer_index] is None:
|
|
return None
|
|
hash_encoder = kvcomp_meta.hash_encoder
|
|
num_tokens = kvcomp_meta.num_actual_tokens
|
|
|
|
hashk = hash_encoder.compute_hash(key[:num_tokens])
|
|
hashk_op = hashk.transpose(0, 1).reshape(-1, hashk.shape[-1]).contiguous()
|
|
|
|
hashk_cache_op = kvcomp_meta.hashk_caches[layer_index]
|
|
real_batch_size = kvcomp_meta.seq_lens_gpu.shape[0]
|
|
|
|
torch.ops._C_ascend.npu_reshape_and_cache_bnsd(
|
|
hashk_op,
|
|
hashk_cache_op,
|
|
kvcomp_meta.slot_mapping[:num_tokens],
|
|
kvcomp_meta.seq_lens_for_reshape[:real_batch_size],
|
|
hashk_cache_op,
|
|
)
|
|
return hashk_cache_op
|
|
|
|
|
|
def get_kvcomp_decode_params(
|
|
layer_index: int | None,
|
|
kvcomp_meta: KVCompMetaData | None,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
block_table: torch.Tensor,
|
|
actual_seq_lengths_kv: list[int],
|
|
):
|
|
assert kvcomp_meta is not None
|
|
assert layer_index is not None
|
|
|
|
if kvcomp_meta.hashk_caches[layer_index] is None:
|
|
return block_table, actual_seq_lengths_kv
|
|
|
|
kv_config = kvcomp_meta.kvcomp_config
|
|
hash_encoder = kvcomp_meta.hash_encoder
|
|
real_batch_size = kvcomp_meta.seq_lens_gpu.shape[0]
|
|
|
|
if kv_config.vllm_hash_attention_skip_layers[layer_index]:
|
|
return kvcomp_meta.hamming_output, kvcomp_meta.seq_lens_from_hamming
|
|
|
|
hashk_cache_op = reshape_and_cache_kvcomp(kvcomp_meta, layer_index, key)
|
|
|
|
hashq = hash_encoder.compute_hash(query[:real_batch_size])
|
|
hashq_op = hashq.unsqueeze(2).contiguous()
|
|
|
|
new_block_table = torch.ops._C_ascend.npu_hamming_dist_top_k(
|
|
hashq_op,
|
|
hashk_cache_op,
|
|
None,
|
|
kvcomp_meta.topk_for_hamming_full[:real_batch_size],
|
|
kvcomp_meta.seq_lens_gpu[:real_batch_size],
|
|
kvcomp_meta.chunk_sizes_for_hamming_full[:real_batch_size],
|
|
kvcomp_meta.max_seq_len_for_hamming,
|
|
kvcomp_meta.sink,
|
|
kvcomp_meta.recent,
|
|
None,
|
|
kvcomp_meta.block_tables_for_hamming,
|
|
kvcomp_meta.valid_query_mask[:real_batch_size],
|
|
kvcomp_meta.hamming_output[:real_batch_size],
|
|
)
|
|
|
|
new_block_table = new_block_table.squeeze(1).contiguous()
|
|
kvcomp_meta.hamming_output = new_block_table
|
|
|
|
return new_block_table, kvcomp_meta.seq_lens_from_hamming
|
|
|
|
|
|
def is_enable_hamming_sparse():
|
|
vllm_config = get_current_vllm_config_or_none()
|
|
if vllm_config is None:
|
|
return False
|
|
additional_config = vllm_config.additional_config if vllm_config.additional_config is not None else {}
|
|
enable_hamming_sparse = additional_config.get("enable_hamming_sparse", False)
|
|
enable_hamming_sparse = enable_hamming_sparse and not vllm_config.speculative_config
|
|
return enable_hamming_sparse
|