init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

vllm_ascend/attention/kvcomp_attn/attention_utils.py

@@ -0,0 +1,152 @@
+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