src/prefill.py

import sys
import os
from typing import List, Dict
from dataclasses import dataclass
import multiprocessing as mp
import pathlib
import torch

from transformers import AutoTokenizer, BitsAndBytesConfig

project_path = pathlib.Path(__file__).parent.parent
sys.path.append(str(project_path))

from src.utils.gpu_worker import GpuWorker
from src.utils.template import QWEN3_TEMPLATE, QWEN3_INSTRUCT_TEMPLATE
from src.utils.cache import  copy_kv_cache_to_device, CustomDynamicCacheOnCPU
from src.msa.model import MSAForCausalLM
from src.utils.gpu_worker import GpuWorker
from src.utils.tools import compose_input
from src.types import Document, ProtocolConstants

@dataclass
class BlockModelInput:
    doc_input_ids: torch.Tensor
    doc_attention_mask: torch.Tensor
    doc_ids: torch.Tensor
    position_ids: torch.Tensor
    num_chunks: int
    chunk_sizes: List[int]

PREFILL_WORKER_READY = "PREFILL_WORKER_READY"
PREFILL_WORKER_CLOSE = "PREFILL_WORKER_CLOSE"
PREFILL_WORKER_MEMORY_DOCS = "PREFILL_WORKER_MEMORY_DOCS"
PREFILL_WORKER_NUM_CHUNKS_REPORT = "PREFILL_WORKER_NUM_CHUNKS_REPORT"
PREFILL_WORKER_META = "PREFILL_WORKER_META"

class PrefillStage1Worker(GpuWorker):
    """Memory工作进程"""

    def __init__(self, gpu_id: int, model_path: str,  template: dict,
                 pooling_kernel_size: int, envs: dict):
        """
        该 worker 被MemoryWorker创建并仅执行prefill stage 1获取block 的kv cache
        """
        super().__init__(gpu_id, envs)
        self.model_path = model_path
        self.pooling_kernel_size = pooling_kernel_size

        self._load_model()
        self.model_config = self.model.config

        self.template_id = -2
        self._prepare_template(template)

    def num_model_layers(self):
        return self.model_config.num_hidden_layers

    def _load_model(self):
        """加载模型和tokenizer"""
        # 加载tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)

        # 加载模型
        self.model = MSAForCausalLM.from_pretrained(
            self.model_path,
            use_cache=True,
            attn_implementation="flash_attention_2",
            torch_dtype="bfloat16",
            device_map=self.device,
        )
        self.model.eval()

    
    @staticmethod
    def split_docs(docs: List[Document], block_size: int):
        sub_blocks = []
        curr_block = []
        sz = 0
        for doc in docs:
            chunks = doc.num_chunks
            if sz + chunks > block_size and curr_block:
                sub_blocks.append(curr_block)
                sz = 0
                curr_block = []
            curr_block.append(doc)
            sz += chunks
        
        if curr_block:
            sub_blocks.append(curr_block)
        
        return sub_blocks
    
    @staticmethod
    def wait_for_ready(q: mp.Queue):               
        ProtocolConstants.expect(q, PREFILL_WORKER_READY)

    @staticmethod
    def close_worker(q: mp.Queue):               
        ProtocolConstants.send(q, PREFILL_WORKER_CLOSE, block=True)

    @staticmethod
    def send_documents(q: mp.Queue, docs):
        ProtocolConstants.send(q,
                               PREFILL_WORKER_MEMORY_DOCS,
                               data=docs,
                               block=False)
    @staticmethod
    def recv_meta(q: mp.Queue):
        return ProtocolConstants.expect(q, PREFILL_WORKER_META)

    @staticmethod
    def prefill_worker_main(gpu_id: int, request_queue: mp.Queue, response_queue: mp.Queue,
                            model_path: str, template: Dict,
                            pooling_kernel_size: int, block_size: int, envs):
        
        # print(f"prefill worker {gpu_id} started")
        worker = PrefillStage1Worker(gpu_id, model_path, template, pooling_kernel_size, envs)

        # notify parent I'm ready
        ProtocolConstants.send(response_queue, PREFILL_WORKER_READY, block=False)

        docs: List[Document] = ProtocolConstants.expect(request_queue, PREFILL_WORKER_MEMORY_DOCS)
        try:
            for block in PrefillStage1Worker.split_docs(docs, block_size):
                meta = worker.inference(block)

                # send to master worker process and continue, DO NOT block
                ProtocolConstants.send(response_queue, 
                                    PREFILL_WORKER_META,
                                    data=meta,
                                    block=False)

        except Exception as e:
            print(f"[子进程 {gpu_id}] 发生错误: {e}")
            import traceback
            traceback.print_exc()
        
        # wait for exit signal
        ProtocolConstants.expect(request_queue, PREFILL_WORKER_CLOSE)
        print(f"prefill worker {gpu_id} ended")

    def inference(self, block: List[Document]):
        """
        处理memory block

        Args:
            memory_block: 分配给此GPU的memory block

        """
        model_input = self._prepare_block_inputs(block)

        kv_meta = self._inference(model_input)
        kv_meta['nr_docs'] = len(block)
        kv_meta['doc_ids'] = [item.doc_id for item in block]
        kv_meta['nr_chunks'] = [item.num_chunks for item in block]

        return kv_meta

    def _prepare_template(self, template: Dict) -> Dict:
        """
        重新加载单个memory block
        完整复制eval_anything_v2_batch.py中reload_memory的逻辑

        Args:
            block: memory block数据 [(doc_id, doc_str), ...]
            template: 模板字典

        Returns:
            KV cache元数据
        """
        # 获取模板信息
        self.pad_token = self.tokenizer.pad_token
        self.pad_token_id = self.tokenizer.pad_token_id
        self.doc_end_id = self.tokenizer("<|im_end|>", add_special_tokens=False)["input_ids"]

        prompt_template = template["prompt"].replace("{prompt}", self.pad_token)
        prompt_template_inputs = self.tokenizer(prompt_template, add_special_tokens=False)
        self.prompt_template_input_ids = prompt_template_inputs["input_ids"]
        self.prompt_template_attention_mask = prompt_template_inputs["attention_mask"]

        self.pad_index = self.prompt_template_input_ids.index(self.pad_token_id)
        self.template_head_input_ids = self.prompt_template_input_ids[:self.pad_index]
        self.template_head_attention_mask = self.prompt_template_attention_mask[:self.pad_index]
        self.template_tail_input_ids = self.prompt_template_input_ids[self.pad_index+1:]
        self.template_tail_attention_mask = self.prompt_template_attention_mask[self.pad_index+1:]


    def _prepare_block_inputs(self, block: List[Document]) -> Dict:
        """
        重新加载单个memory block
        完整复制eval_anything_v2_batch.py中reload_memory的逻辑

        Args:
            block: memory block数据 [(doc_id, doc_str), ...]
            template: 模板字典

        Returns:
            KV cache元数据
        """

        # 准备文档数据

        # block starts with template head
        doc_ids = [self.template_id] * len(self.template_head_input_ids)
        doc_input_ids = [i for i in self.template_head_input_ids ]
        doc_attention_mask = [i for i in self.template_head_attention_mask]
        position_ids = [i for i in range(self.pad_index)]

        chunk_sizes = [] # 记录每一份文档占用的 chunk 数量

        for doc_idx, item in enumerate(block):
            doc_id, doc, pre_calculated_num_chunk = item.doc_id, item.doc, item.num_chunks
            new_doc, doc_inputs = compose_input(doc, doc_id, self.tokenizer)
            # print(f"inference {doc_id+1}: {new_doc}")
            # 此处必须使用 1 起始的doc_idx，因为此 id 是用于生成 pool doc ID 的
            # 注意不可以使用doc_id，doc_id只能用于嵌入在语料中，使得 generate 时能生成出来，
            # 而pool doc ID的作用却是用于标注产生的 kv cache chunks 和文档的对应关系
            # 所以每次 stage1 的推理doc id 都是从 1 开始的
            temp_doc_ids = [doc_idx+1] * len(doc_inputs["input_ids"])
            temp_doc_input_ids = doc_inputs["input_ids"]
            temp_doc_attention_mask = doc_inputs["attention_mask"]
            length = len(temp_doc_input_ids)
            temp_position_ids = [i for i in range(length)]

            chunk_size = (len(temp_doc_ids) + self.pooling_kernel_size - 1) // self.pooling_kernel_size
            chunk_sizes.append(chunk_size)
            assert chunk_size == pre_calculated_num_chunk, f"pre calculated chunk {pre_calculated_num_chunk} got {chunk_size}, doc str {len(doc)} id len {length}/{len(temp_doc_ids)}: [{doc_id}] <{doc}>"



            doc_ids.extend(temp_doc_ids)
            doc_input_ids.extend(temp_doc_input_ids)
            doc_attention_mask.extend(temp_doc_attention_mask)
            position_ids.extend(temp_position_ids)

        input_ids_tensor = torch.LongTensor([doc_input_ids])
        attention_mask_tensor = torch.LongTensor([doc_attention_mask])
        doc_ids_tensor = torch.LongTensor([doc_ids])
        position_ids_tensor = torch.LongTensor([position_ids])

        return BlockModelInput(doc_input_ids=input_ids_tensor,
                               doc_attention_mask=attention_mask_tensor,
                               doc_ids=doc_ids_tensor,
                               position_ids=position_ids_tensor,
                               num_chunks=sum(chunk_sizes),
                               chunk_sizes=chunk_sizes)

    def _inference(self, model_input: BlockModelInput) -> Dict:
        """
        重新加载单个memory block
        完整复制eval_anything_v2_batch.py中reload_memory的逻辑

        Args:
            block: memory block数据 [(doc_id, doc_str), ...]
            template: 模板字典

        Returns:
            KV cache元数据
        """

        # 转换为tensor
        input_ids_tensor = model_input.doc_input_ids.to(self.device)
        attention_mask_tensor = model_input.doc_attention_mask.to(self.device)
        doc_ids_tensor = model_input.doc_ids.to(self.device)
        position_ids_tensor = model_input.position_ids.to(self.device)

        # 创建past_key_values，这里我们会直接将 kvcache 和其他的 tensor 全部保存在 cpu 上
        past_key_values = CustomDynamicCacheOnCPU()
        for layer_idx in range(self.num_model_layers()):
            past_key_values.record_kwargs(layer_idx, {"stage": "prefill_stage1"})

        # 执行prefill
        # TODO: 多 batch 会更快
        with torch.no_grad():
            if True:
                """我们的数据太大了，会导致model lm_head产生大量的显存堆积，所以直接用model.model避过去"""
                outputs = self.model.model(
                    input_ids=input_ids_tensor,
                    attention_mask=attention_mask_tensor,
                    position_ids=position_ids_tensor,
                    past_key_values=past_key_values,
                    use_cache=True,
                    output_attentions=False,
                    output_hidden_states=False,
                    output_docs_score=False,
                    doc_ids=doc_ids_tensor,
                )
            else:
                outputs = self.model(
                    input_ids=input_ids_tensor,
                    attention_mask=attention_mask_tensor,
                    doc_ids=doc_ids_tensor,
                    use_cache=True,
                    position_ids=position_ids_tensor,
                    past_key_values=past_key_values,
                )

        torch.cuda.empty_cache()
        # 构建返回的元数据
        kvcache_meta = {
            "chunk_sizes": model_input.chunk_sizes,
            "past_key_values": outputs.past_key_values,
        }

        return kvcache_meta