src/msa_service.py

import queue
import pickle
import json
from tqdm import tqdm
import torch
import torch.distributed as dist
import multiprocessing as mp
import threading
import os
from typing import Tuple, List, Optional, Dict, Union, Any, Callable
from dataclasses import dataclass
from transformers import AutoTokenizer, BitsAndBytesConfig, AutoConfig

from src.config.memory_config import GenerateConfig, MemoryConfig, ModelConfig
from src.prefill import PrefillStage1Worker
from src.types import ProtocolConstants
from src.utils.cache import create_cache, CustomDynamicCache
from src.utils.gpu_worker import GpuWorker
from src.utils.tools import RequestLimiter, format_bytes, cumulative_concat, compose_input
from src.msa.model import MSAForCausalLM, MSAConfig

# ==========================================
# 数据结构定义
# ==========================================

MEMORY_WORKER_READY = "MEMORY_WORKER_READY"
MEMORY_WORKER_CLOSE = "MEMORY_WORKER_CLOSE"
MEMORY_WORKER_BLOCKS = "MEMORY_WORKER_BLOCKS"
MEMORY_WORKER_IDX_TO_DOC = "MEMORY_WORKER_IDX_TO_DOC"
MEMORY_WORKER_REPORT_KV_STATES = "MEMORY_WORKER_REPORT_KV_STATES"

@dataclass
class CmdBase:
    name: str = ""

@dataclass
class ResultBase:
    name: str = ""
    gpu_id: int = 0


@dataclass
class QueryTemplateKVPrefixCmd(CmdBase):
    layer_idx: int = -1 # -1 indicate get all layers
    
    def __post_init__(self):
        self.name = "query_template_kv_prefix"

@dataclass
class QueryTemplateKVPrefixResult(ResultBase):
    template_kvcache: Dict[int, Dict[str, torch.Tensor]] = None
    
    def __post_init__(self):
        self.name = "query_template_kv_prefix"

@dataclass
class PrefillStage2Cmd(CmdBase):
    query_states: torch.Tensor = None
    attention_mask: torch.Tensor = None
    doc_ids: torch.Tensor = None
    layer_index: int = 0
    
    def __post_init__(self):
        self.name = "prefill_stage2"

@dataclass
class PrefillStage2Result(ResultBase):
    key: Union[torch.Tensor, List[torch.Tensor]] = None
    value: Union[torch.Tensor, List[torch.Tensor]] = None
    pooled_doc_ids: Union[torch.Tensor, List[torch.Tensor]] = None
    
    def __post_init__(self):
        self.name = "prefill_stage2"

# TODO: 这个可以取消，因为这是为了让prefill产生的 pooled doc id 全局唯一，但是这个完全可以
#       在归集stage2结果的时候才去做，只要用 block 或者 gpu id 做一个偏移就好
@dataclass
class SetDocIDCmd(CmdBase):
    pooled_doc_id_bias: int = 0
    
    def __post_init__(self):
        self.name = "set_doc_id"

@dataclass
class ReportDocID(ResultBase):
    pooled_doc_id_bias: int = 0
    
    def __post_init__(self):
        self.name = "report_doc_id"

class CustomDynamicCacheOnCPU(CustomDynamicCache):
    def __init__(self, _distributed_cache_data=None):
        super().__init__(_distributed_cache_data)

    # def record_kwargs(self, layer_idx, kwargs):
    #     d = {}
    #     for k, v in kwargs.items():
    #         if v is not None and torch.is_tensor(v):
    #             d[k] = v.cpu() if v.is_cuda else v.clone()
    #         else:
    #             d[k] = v
    #     super().record_kwargs(layer_idx, kwargs)

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs=None,
        ) -> tuple[torch.Tensor, torch.Tensor]:
        # if key_states is not None and torch.is_tensor(key_states) and key_states.is_cuda:
        #     key_states = key_states.cpu()
        if value_states is not None and torch.is_tensor(value_states) and value_states.is_cuda:
            value_states = value_states.cpu()
        return super().update(key_states, value_states, layer_idx, cache_kwargs)

@dataclass
class GenerateRequest(CmdBase):
    # 每次用户发送一个 generate 请求，我们都分配一个 msgid
    # 但是我们会将其分成world个GenerateRequest，seq_id表示其自请求顺序
    msg_id: int = 0
    seq_id: int = 0
    input_ids: List[int] = None
    attention_mask: List[int] = None
    doc_ids: List[int] = None
    positions: List[int] = None
    require_recall_topk: bool = False
    prompts: List[str] = None
    
    def __post_init__(self):
        self.name = "generate_request"

@dataclass
class GenerateResponse():
    msg_id: int = -1
    seq_id: int = -1
    gpu_id: int = 0
    generated_texts: List[str] = None
    recall_topk: Dict[int, List] = None


# response receiver: callback(generated_texts, required_recall_topk, userdata)
GenerateReceiver = Callable[[List[str], Any, Any], None]

@dataclass
class GenerateStub:
    msg_id: int = 0
    userdata: Any = None
    responses: List[GenerateResponse] = None
    nr_dummy: int = 0 # 标记这个 batch 尾部有几个query 是填充的空请求
    callback: Optional[GenerateReceiver] = None

    def respond(self):
        if self.callback is None:
            print("no callback")
            return

        self.responses.sort(key=lambda r: r.seq_id)
        txts = [] 
        recall_topk: Dict[int, List] = {}
        for rsp in self.responses:
            txts.extend(rsp.generated_texts)
            if rsp.recall_topk:
                for layer, topks in rsp.recall_topk.items():
                    if layer not in recall_topk:
                        recall_topk[layer] = []
                    recall_topk[layer].extend(topks)
        
        self.callback(txts[:len(txts)-self.nr_dummy], recall_topk, self.userdata)
            
                
            
        

@dataclass
class Document:
    doc: str = ""
    doc_id: int = 0
    num_chunks: int = 0

@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]

@dataclass
class BlockDesc:
    """一个 block 的描述数据"""
    docs: List[Document] = None # 所有分配到此 GPU 上的文档
    tmp_doc_ids: List[torch.Tensor] = None # 临时存储

    # kvcache的每一个 chunk 分组编号，从 1 开始，形如[1,1,1,2,2,3,3,3,...]
    # pool_ids[-1] 就是文档总数，相同的编号对应的 chunks 表示同属于同一个文档
    pool_ids: torch.Tensor = None
    # 用于通过 pool_id查找 global doc id，即 global_doc_ids = doc_ids[pool_ids]
    doc_ids: torch.Tensor = None
    doc_ids_cpu: torch.Tensor = None # doc_ids 的 CPU 版本，方便后续索引使用

    # 用 local id，也就是pool_ids中的 id 来进行检索
    doc_offsets_cpu: torch.Tensor = None # 每个文档 chunk 的起始索引 [nr_docs + 1]
    doc_lens_cpu: torch.Tensor = None    # 每个文档包含的 chunk 数量 [nr_docs + 1]

    def init_docs(self, docs: List[Document], device):
        self.docs = docs
        self.doc_ids_cpu = torch.LongTensor([0] + [item.doc_id for item in self.docs])
        self.doc_ids = self.doc_ids_cpu.to(device)
        self.tmp_doc_ids = []

        # 计算偏移量：假设 doc 0 长度为 0
        lens = [0] + [doc.num_chunks for doc in docs]
        self.doc_lens_cpu = torch.LongTensor(lens)
        self.doc_offsets_cpu = torch.cumsum(torch.LongTensor([0] + lens[:-1]), dim=0)

    def create_global_doc_ids(self, local_doc_ids):
        return self.doc_ids[local_doc_ids]

    def merge_poolig_doc_id(self, device):
        self.pool_ids_cpu = cumulative_concat(self.tmp_doc_ids)
        self.pool_ids = self.pool_ids_cpu.to(device)
        self.tmp_doc_ids.clear()

    def chunks(self):
        return sum(doc.num_chunks for doc in self.docs)
    
    @property
    def nr_docs(self):
        return len(self.docs)

@dataclass
class BlockData:
    """本 GPU 上的文档 kvcache：[1, num_head, num_chunks, head_dim]"""
    k: torch.Tensor = None # k 用于 attn 计算，如果 rk 不存在则保存于 GPU，同时用于检索，否则保存在 CPU
    v: torch.Tensor = None # 存储于 CPU
    rk: torch.Tensor = None # 用于检索，此时我们有独立的检索 k，需要保存于 GPU

    def get_router_k(self):
        """获取用于检索的k"""
        return self.rk if self.rk is not None else self.k
    

@dataclass
class SliceDesc:
    """block分片信息，每一个分片包含了一部分文档信息
    """
    nr_docs: int = 0 # 文档数量
    nr_chunks: int = 0 # chunk数量
    global_doc_ids: torch.Tensor = None # 文档 ID，shape 为[nr_chunks]

    # 下面这几个 tensor 是我们预先计算好，用于 prefill stage2 的数据
    local_doc_ids: torch.Tensor = None
    local_doc_ids_0: torch.Tensor = None
    original_doc_ids: torch.Tensor = None


class MemoryClientBase():
    """被模型调用的接口"""
    def get_template_prefix_kvcaches(self, layer_idx: int):
        pass
    def doc_query(self, query_states: torch.Tensor, attention_mask: torch.Tensor, layer_idx: int):
        pass

# ==========================================
# Service 实现
# ==========================================
class Memory(GpuWorker):
    """Memory工作进程"""

    def __init__(self,
                 gpu_id: int,
                 generate_config: GenerateConfig,
                 model_config: ModelConfig,
                 memory_config: MemoryConfig):
        """
        该 worker 被MemoryWorker创建并仅执行prefill stage 1获取block 的kv cache
        """
        super().__init__(gpu_id, model_config.get_model_envs())

        self.model_config = model_config
        self.generate_config = generate_config
        self.memory_config = memory_config

        self.msa_model_config = MSAConfig.from_pretrained(model_config.model_path)

        self.router_layer_ids = list(range(self.num_model_layers()) if model_config.router_layer_idx == "all"
                                     else map(int, model_config.router_layer_idx.split(",")))
        self.router_layer_ids.sort()

        self.num_key_value_groups = self.msa_model_config.num_attention_heads // self.msa_model_config.num_key_value_heads


        self.vcache_in_cpu = True
        self.v_device = 'cpu' if self.vcache_in_cpu else self.device

        ######## 下面几个定义是可以动态加载进来的
        # 每一个router层上的分片数据列表: layer_idx => [block 1, block 2, ...]
        self.blocks: Dict[int, BlockData] = {layer: BlockData() for layer in self.router_layer_ids}
        # 每一个 block 的描述数据, [block 1 desc, block 2 desc, ...]
        self.block_desc: BlockDesc = BlockDesc()
        # 每一层上的template prefix kv cache:  layer idx -> (k, v)
        self.template_prefix_kvcache: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {}

        # 对 block 进行chunk 分割时，每个 block 大小
        self.slice_chunk_size = memory_config.slice_chunk_size
        # 为了后续计算方便，我们预先处理好了 k slice，shape 为[1, num_kv_heads, 1, hdim, chunk]
        self.k_slices: Dict[int, List[torch.Tensor]] = []
        self.slice_desc: List[SliceDesc] = []

        self.head_reduce_method = self.msa_model_config.msa_config.head_reduce_method
        self.query_reduce_method = self.msa_model_config.msa_config.query_reduce_method
        self.chunk_reduce_method = self.msa_model_config.msa_config.chunk_reduce_method
        self.aux_loss_method = self.msa_model_config.msa_config.aux_loss_method
        self.decouple_router = self.msa_model_config.msa_config.decouple_router

        if self.aux_loss_method == "INFONCE" and self.decouple_router:
            self.scaling = -1.0
        else:
            self.scaling = self.msa_model_config.head_dim**-0.5

    def _start_worker(self):
        self._worker_req_q, self._worker_rsp_q = mp.Queue(), mp.Queue()

        self._worker_process = mp.Process(
            target=PrefillStage1Worker.prefill_worker_main,
            args=(self.gpu_id, self._worker_req_q, self._worker_rsp_q,  
                  self.model_config.model_path,  self.generate_config.template, 
                  self.memory_config.pooling_kernel_size, self.memory_config.block_size,
                  self.model_config.get_model_envs()),
            name=f"PrefillStage1Worker-{self.gpu_id}"
        )
        self._worker_process.start()

        PrefillStage1Worker.wait_for_ready(self._worker_rsp_q)
        print(f"Prefill worker {self.gpu_id} is ready")
    
    def _stop_worker(self):
        print(f"Stop prefill worker on GPU {self.gpu_id}")
        PrefillStage1Worker.close_worker(self._worker_req_q)
        self._worker_process.join()
        self._worker_process = None
        self._worker_req_q, self._worker_rsp_q = None, None


    def num_model_layers(self):
        return self.msa_model_config.num_hidden_layers
    
    
    def serialize(self, path: str):
        """
        安全分片序列化：
        1. 创建目录
        2. 保存元数据 (meta.pt)
        3. 逐层保存 Tensor 到独立文件，避免内存峰值
        """
        # 1. 准备目录 (如果存在则警告或清空，这里选择清空以确保一致性)
        if os.path.exists(path):
            print(f"Warning: Directory {path} exists. Overwriting data inside.")
        else:
            os.makedirs(path, exist_ok=True)
        
        stub = {
            "memory_file_path": self.memory_config.memory_file_path,
            "world": self.generate_config.world,
        }
        with open(os.path.join(path, "stub.json"), "w") as fp:
            json.dump(stub, fp)

        # 2. 保存元数据 (BlockDesc + 结构信息)
        # 注意：BlockData 里的 heavy tensor 不在这里存，只存结构
        meta_dict = {
            "router_layer_ids": self.router_layer_ids,
            "block_desc": {
                "docs": self.block_desc.docs,
                "doc_ids": self.block_desc.doc_ids.cpu(),
                # pool_ids 转 CPU 保存
                "pool_ids": self.block_desc.pool_ids.cpu() if self.block_desc.pool_ids is not None else None
            },
        }
        torch.save(meta_dict, os.path.join(path, "meta.pt"))

        # 3. 逐层保存 Blocks (重头戏)
        print("Serializing Blocks (Sharded)...")
        for layer_id, block in tqdm(self.blocks.items(), desc="Saving Layers"):
            # 分离存储 K 和 V
            if block.k is not None:
                # K 是 GPU -> CPU
                torch.save(block.k.cpu(), os.path.join(path, f"layer_{layer_id}_k.pt"))
            
            if block.v is not None:
                # V 已经在 CPU，直接存
                # 提示：对于 CPU 上的极大 Tensor，torch.save 效率很高
                torch.save(block.v.cpu(), os.path.join(path, f"layer_{layer_id}_v.pt"))

            if block.rk is not None:
                torch.save(block.rk.cpu(), os.path.join(path, f"layer_{layer_id}_rk.pt"))

        # 4. 保存 Prefix Cache
        if self.template_prefix_kvcache:
            print("Serializing Prefix Cache...")
            prefix_data = {}
            for layer_id, (k, v) in self.template_prefix_kvcache.items():
                prefix_data[layer_id] = (k.cpu(), v.cpu())
            torch.save(prefix_data, os.path.join(path, "prefix_cache.pt"))

        print(f"Serialization complete. Data saved to {path}")

    def deserialize(self, path: str):
        """
        流式反序列化：
        每次只加载一个文件进内存，处理完后立即释放或转移到 GPU，
        将系统内存（RAM）占用控制在最低。
        """
        if not os.path.exists(path):
            raise FileNotFoundError(f"Path {path} does not exist.")
        
        print("Deserialize from ", path)

        # 1. 加载元数据
        meta_path = os.path.join(path, "meta.pt")
        meta = torch.load(meta_path, map_location="cpu", weights_only=False)
        
        # 恢复 BlockDesc
        bd_info = meta["block_desc"]
        self.block_desc = BlockDesc(
            docs=bd_info["docs"],
            doc_ids=bd_info["doc_ids"].to(self.device),
            pool_ids=bd_info["pool_ids"].to(self.device) if bd_info["pool_ids"] is not None else None
        )

        # 2. 逐层恢复 Blocks
        # 关键：我们不一次性加载所有文件，而是循环加载
        self.blocks = {}
        print(f"Deserializing Blocks to GPU {self.gpu_id}...")
        
        for layer_id in tqdm(self.router_layer_ids, desc=f"GPU-{self.gpu_id} Loading Layers", position=self.gpu_id):
            k_path = os.path.join(path, f"layer_{layer_id}_k.pt")
            v_path = os.path.join(path, f"layer_{layer_id}_v.pt")
            rk_path = os.path.join(path, f"layer_{layer_id}_rk.pt")
            
            current_k = None
            current_v = None
            current_rk = None

            # 加载 K -> 立即转 GPU -> 删除 CPU 副本
            if os.path.exists(k_path):
                temp_k = torch.load(k_path, map_location="cpu", weights_only=False)
                current_k = temp_k.to(self.device) # 移入显存
                del temp_k # 立即释放内存

            # 加载 V -> 驻留 CPU
            if os.path.exists(v_path):
                # 既然 V 是 CPU 存储且极大，直接 load 即可
                # 如果 V 极其巨大（比如超过 RAM），可以考虑 mmap=True，
                # 但这要求保存时格式兼容，且 V 通常作为 cache 需要频繁读取，RAM 还是最快的
                current_v = torch.load(v_path, map_location="cpu", weights_only=False)

            if os.path.exists(rk_path):
                current_rk = torch.load(rk_path, map_location="cpu", weights_only=False)

            self.blocks[layer_id] = BlockData(k=current_k, v=current_v, rk=current_rk)
            
            # 显式触发垃圾回收（可选，但在内存吃紧时有用）
            # import gc; gc.collect() 

        # 3. 恢复 Prefix Cache
        prefix_path = os.path.join(path, "prefix_cache.pt")
        if os.path.exists(prefix_path):
            prefix_raw = torch.load(prefix_path, map_location="cpu", weights_only=False)
            self.template_prefix_kvcache = {}
            for layer_id, (k_cpu, v_cpu) in prefix_raw.items():
                self.template_prefix_kvcache[layer_id] = (
                    k_cpu.to(self.device),
                    v_cpu.to(self.device)
                )
        
        print("Deserialization complete.")
    
    def _init_block_data(self, shape: torch.Size, dtype: torch.dtype, has_rk=False):
        total_chunks = self.block_desc.chunks()
        for layer_idx in self.router_layer_ids:
            block_data = self.blocks[layer_idx]

            if block_data.k is None:
                kv_shape = list(shape)
                kv_shape[2] = total_chunks
                # TODO: deserialize 时需要考虑 pin_memory
                block_data.k = torch.empty(kv_shape, device=self.v_device if has_rk else self.device, dtype=dtype, pin_memory=has_rk)
                block_data.v = torch.empty(kv_shape, device=self.v_device, dtype=dtype, pin_memory=True)
                block_data.rk = torch.empty(kv_shape, device=self.device, dtype=dtype) if has_rk else None

    def save_idx_to_doc(self, idx_to_doc: Dict[int, str]):
        self.idx_to_doc = idx_to_doc
        
    def generate_blocks(self, docs: List[Document]):
        memory_data_path = os.environ.get("MEMORY_DATA_PATH", "")
        if memory_data_path:
            memory_data_path = os.path.join(memory_data_path, f"gpu{self.generate_config.world}", str(self.gpu_id))
            if os.path.exists(memory_data_path):
                self.deserialize(memory_data_path)
                assert len(docs) == self.block_desc.nr_docs
                self._post_process()
                return

        self._start_worker()

        self.block_desc.init_docs(docs, self.device)

        PrefillStage1Worker.send_documents(self._worker_req_q, docs)

        kv_offset = 0
        recv = 0

        total_docs = len(docs)
        total_chunks = self.block_desc.chunks()
        pbar = tqdm(total=total_docs, desc=f"Worker-{self.gpu_id} memory inference", position=self.gpu_id)
        while recv < total_docs:
            meta = PrefillStage1Worker.recv_meta(self._worker_rsp_q)

            recv += meta['nr_docs'] # 此 meta 包含了几个文档
            pbar.update(meta['nr_docs'])

            past_key_values: CustomDynamicCacheOnCPU = meta["past_key_values"]

            # 仅保存一次
            if not self.template_prefix_kvcache:
                kwargs = past_key_values.cache_kwargs
                self.template_prefix_kvcache = {
                    layer_idx: (
                        kwargs[layer_idx]['template_prefix_kcache'].to(self.device),
                        kwargs[layer_idx]['template_prefix_vcache'].to(self.device)
                        )
                    for layer_idx in range(self.num_model_layers())
                }


            meta_chunks = 0 # 这次的 meta 有多少 chunks

            for layer_idx in self.router_layer_ids:
                block_data = self.blocks[layer_idx]

                k, v = past_key_values.get_kvcache(layer_idx) # for attn calc
                rk = past_key_values.get_router_kcache(layer_idx) # for router, if any

                desc: BlockDesc = self.block_desc
                if block_data.k is None:
                    self._init_block_data(k.shape, k.dtype, has_rk=rk is not None)

                if layer_idx == self.router_layer_ids[0]:
                    # 聚合 pooling doc id，因为每一层上的 pooling doc id 都是一样的，所以只聚合第一层
                    desc.tmp_doc_ids.append(past_key_values.cache_kwargs[layer_idx]['pooled_doc_ids'])
                    
                k_cache, v_cache, rk_cache = block_data.k, block_data.v, block_data.rk
                if meta_chunks == 0:
                    meta_chunks = k.shape[2]
                else:
                    assert meta_chunks == k.shape[2], f"expect meta_chunks {meta_chunks} but got {k.shape[2]}"
                    assert kv_offset + meta_chunks <= total_chunks, f"overflow: offset {kv_offset}, meta_chunks {meta_chunks}, total_chunks {total_chunks}"

                # print(f"shape {k.shape} offset {kv_offset}")
                k_cache[:, :, kv_offset:kv_offset+meta_chunks] = k.to(k_cache.device)
                v_cache[:, :, kv_offset:kv_offset+meta_chunks] = v.to(v_cache.device)
                if rk is not None:
                    rk_cache[:, :, kv_offset:kv_offset+meta_chunks] = rk.to(rk_cache.device)

            # 每次一个 meta 处理完才需要进偏移
            kv_offset += meta_chunks

            past_key_values.clear_kvcache()

        pbar.close()
        self._stop_worker()

        # 合并临时pooled doc id
        self.block_desc.merge_poolig_doc_id(self.device)

        if memory_data_path:
            self.serialize(memory_data_path)

        self._post_process()

    def _post_process(self):
        # 产生切片
        self._generate_slice(self.slice_chunk_size)

        kv_bytes = self.blocks[self.router_layer_ids[0]].k.nbytes * len(self.router_layer_ids)
        if not self.vcache_in_cpu:
            kv_bytes *= 2
        
        kv_shape = self.blocks[self.router_layer_ids[0]].k.shape

        print(f"GPU {self.gpu_id} memory usage: {format_bytes(kv_bytes)} kv shape {kv_shape}")
        return kv_bytes, kv_shape

    @staticmethod
    def map_tensor_to_group_ids(a: torch.Tensor) -> torch.Tensor:
        """
        根据相邻元素是否相同，将输入 Tensor a 的值映射到递增的组 ID Tensor b。

        a[i] == a[i-1] => b[i] = b[i-1]
        a[i] != a[i-1] => b[i] = b[i-1] + 1
        b[0] 始终为 1。

        Args:
            a (torch.Tensor): 待处理的一维 int Tensor，必须在 GPU 上。

        Returns:
            torch.Tensor: 生成的组 ID Tensor b，在 GPU 上。
        """
        if a.ndim != 1:
            raise ValueError("输入 Tensor a 必须是一维的。")

        # 1. 检查相邻元素是否不同 (a[i] != a[i-1])
        # torch.diff(a) 计算相邻元素之间的差值 a[i] - a[i-1]
        # 差值不为 0 (!= 0) 即表示 a[i] != a[i-1]
        # 结果是一个布尔 Tensor，长度比 a 短 1，代表从 a[1] 开始的相邻比较
        diff_mask = torch.diff(a) != 0  # [L-1]

        # 2. 将布尔值转换为整数 (0 或 1)
        # True -> 1 (需要增加 ID)，False -> 0 (保持 ID)
        # torch.diff 的结果长度是 L-1，我们只关心从 a[1] 开始的递增
        id_increments = diff_mask.int()  # [L-1]

        # 3. 累计增量
        # torch.cumsum 对增量进行累加：
        # 结果: [0, 0+1, 0+1+0, ...]
        # 长度仍然是 L-1
        group_indices_offset = torch.cumsum(id_increments, dim=0)  # [L-1]

        # 4. 构造最终的 Tensor b
        # b[0] 始终是 1
        # 对于 i > 0，b[i] = group_indices_offset[i-1] + 1
        
        # 在 group_indices_offset 前面添加一个 0 作为 a[0] 的基准偏移
        # [0, group_indices_offset[0], group_indices_offset[1], ...]
        # 结果长度为 L
        b = torch.cat((
            torch.tensor([0], device=a.device, dtype=a.dtype), 
            group_indices_offset
        )) + 1
        return b
    def _generate_slice(self, chunks):
        nr_doc_slices = []  # 每个 block slice 所含的文档数量
        nr_chunk_slices = [] # 每个 block slice 所含的 chunk 数量
        chunk_slices_offset = [] # 每个 block slice 所含的 chunk 位置
        chunk_slice_position: Tuple[int, int] = [] # 每个 chunk slice 的起始和结束位置

        current_chunks = 0
        current_docs = 0
        chunk_offset = 0
        for doc in self.block_desc.docs:
            current_chunks += doc.num_chunks
            current_docs += 1
            chunk_offset += doc.num_chunks
            
            if current_chunks >= chunks:
                nr_doc_slices.append(current_docs)
                nr_chunk_slices.append(current_chunks)
                chunk_slices_offset.append(chunk_offset)
                current_chunks = 0
                current_docs = 0

        if current_chunks > 0:
            if current_chunks < 1024 and len(nr_doc_slices) > 0:
                # too small, append to last
                nr_doc_slices[-1] += current_docs
                nr_chunk_slices[-1] += current_chunks
                chunk_slices_offset[-1] = chunk_offset
            else:
                nr_doc_slices.append(current_docs)
                nr_chunk_slices.append(current_chunks)
                chunk_slices_offset.append(chunk_offset)

        starts = [0] + chunk_slices_offset[:-1]
        chunk_slice_position = [(start, end) for start, end in zip(starts, chunk_slices_offset)]

        self.k_slices = {}
        self.slice_desc = []
        for slice_idx, (start, end) in enumerate(chunk_slice_position):
            desc = SliceDesc(nr_docs=nr_doc_slices[slice_idx],
                             nr_chunks=nr_chunk_slices[slice_idx],
                             global_doc_ids=self.block_desc.pool_ids[start:end],
                             local_doc_ids=Memory.map_tensor_to_group_ids(self.block_desc.pool_ids[start:end]))
            desc.local_doc_ids_0 = (desc.local_doc_ids - 1).view(1, -1)
            unique_indices = torch.cat([torch.tensor([0], device=self.device), torch.where(desc.local_doc_ids[1:] != desc.local_doc_ids[:-1])[0] + 1])
            desc.original_doc_ids = desc.global_doc_ids[unique_indices].view(1, -1)
            
            self.slice_desc.append(desc)
        
        for layer_idx in self.router_layer_ids:
            rk = self.blocks[layer_idx].get_router_k()
            slice = [rk[:,:,start:end].transpose(-1, -2).unsqueeze(2) for start, end in chunk_slice_position]
            self.k_slices[layer_idx] = slice

    def get_max_local_pool_doc_id(self) -> List[int]:
        return self.block_desc.pool_ids[-1].item()
        
    def get_template_prefix_kvcaches(self, layer_idx: int):
        return self.template_prefix_kvcache[layer_idx]

    def prefill_stage2(self, layer_idx: int, query_states: torch.Tensor, query_mask: torch.Tensor):
        # query_mask: [bsz, 1, 1, seqlen, 1]
        # query_states: [bsz, n_kv_head, n_kv_groups, seqlen, hdim]

        bsz, num_kv_heads, num_kv_groups, seq_len, head_dim = query_states.shape
        min_val = torch.finfo(query_states.dtype).min
        routing_mask = ~query_mask
        query_mask = query_mask.squeeze(2)  # [bsz, 1, seqlen, 1]
        q_lens = None

        # 初始化分数表
        # 严重注意：请确认 nr_docs 是否真的等于 max_doc_id + 1。
        # 如果 pool_ids 是全局唯一的（例如加了 bias 100000），这里的 total_docs 必须能容纳最大的 ID。
        total_docs = self.block_desc.nr_docs + 1 
        global_doc_scores = torch.full((bsz, total_docs), min_val, dtype=query_states.dtype, device=self.device)

        # --- Phase 1: 打分 (移除 Query Batch Loop) ---
        
        # 针对 A100，如果 bsz < 256，直接全量计算收益最高
        # 我们直接遍历 slice，不再对 query 进行切片
        
        for slice_desc, k_slice_t in zip(self.slice_desc, self.k_slices[layer_idx]):
            # k_slice_t: [1, num_kv_heads, 1, hdim, chunk] (GPU)
            
            # 1. Matmul (Compute Bound)
            # [bsz, num_kv_heads, groups, seqlen, hdim] @ [1, num_kv_heads, 1, hdim, chunk] 
            # -> [bsz, num_kv_heads, groups, seqlen, chunk]
            attn_scores = torch.matmul(query_states, k_slice_t)
            
            # 2. Mask & Reduce (Memory Bound)
            attn_scores.masked_fill_(routing_mask, min_val)
            # Flatten & Reduce heads/seqlen
            # flatten(1,2): [bsz, num_kv_heads*groups, seqlen, chunk]
            # max(dim=2)[0]: [bsz, num_kv_heads*groups, chunk]
            # max(dim=1)[0]: [bsz, chunk]
            # max_scores_per_chunk = attn_scores.flatten(1, 2).max(dim=2)[0].max(dim=1)[0] # [bsz, chunk]
            max_scores_per_chunk = attn_scores.flatten(1, 2)

            # [bsz, num_kv_heads*groups, chunk]
            if self.head_reduce_method == "max":
                max_scores_per_chunk = max_scores_per_chunk.max(dim=1).values
            elif self.head_reduce_method == "mean":
                max_scores_per_chunk = max_scores_per_chunk.mean(dim=1)
            else:
                raise NotImplementedError(f"Unsupported head reduce method: {self.head_reduce_method}")
            
            if self.query_reduce_method == "max":
                max_scores_per_chunk = max_scores_per_chunk.max(dim=1).values
            elif self.query_reduce_method == "mean":
                # 1. 调整 mask 形状以匹配 max_scores_per_chunk
                # [bsz, 1, seqlen, 1] -> [bsz, seqlen, 1]
                mask = query_mask.squeeze(1) 

                # 2. 在求和前，将无效位置（min_val）归零
                # 使用 torch.where 比乘法更安全，且能有效处理极小的 min_val
                masked_sum = torch.where(mask, max_scores_per_chunk, torch.zeros_like(max_scores_per_chunk)).sum(dim=1)

                # 3. 计算每个 sequence 中的有效 token 数量
                # 防止除以 0，使用 clamp 或加一个极小值 eps
                valid_counts = mask.sum(dim=1).clamp(min=1.0) 

                # 4. 得到最终均值 [bsz, chunk]
                max_scores_per_chunk = masked_sum / valid_counts
            elif self.query_reduce_method == "last":
                if q_lens is None:
                    q_lens = query_mask.squeeze(3).squeeze(1).sum(dim=1).long()
                
                # 边界保护: 如果 length 为 0 (异常情况), clamp 到 0
                last_indices = (q_lens - 1).clamp(min=0)
                
                # 3. 调整形状用于 gather
                # [B, 1, C]
                gather_idx = last_indices.view(bsz, 1, 1).expand(-1, 1, slice_desc.nr_chunks)
                
                # 4. 提取
                max_scores_per_chunk = max_scores_per_chunk.gather(1, gather_idx).squeeze(1)
            else:
                raise NotImplementedError(f"Unsupported query reduce method: {self.query_reduce_method}")

            # 3. Local Scatter (Chunk -> Doc)
            # slice_desc.local_doc_ids_0 已经在 GPU 上预计算好了
            # [1, chunk] -> [bsz, chunk]
            scatter_indices = slice_desc.local_doc_ids_0.expand(bsz, -1)
            
            local_doc_scores = torch.full((bsz, slice_desc.nr_docs), -float('inf'), device=self.device, dtype=query_states.dtype)
            if self.chunk_reduce_method == "max":
                local_doc_scores.scatter_reduce_(
                    dim=1, index=scatter_indices, src=max_scores_per_chunk, 
                    reduce="amax", include_self=True
                )
            
            elif self.chunk_reduce_method == "mean":
                doc_sums = torch.zeros_like(local_doc_scores)
                doc_sums = doc_sums.scatter_reduce(
                    dim=1, index=scatter_indices, src=max_scores_per_chunk, 
                    reduce="sum", include_self=False
                )
                doc_counts = torch.zeros_like(local_doc_scores)
                ones = torch.ones_like(max_scores_per_chunk)
                doc_counts = doc_counts.scatter_reduce(
                    dim=1, index=scatter_indices, src=ones, 
                    reduce="sum", include_self=False
                )
                doc_counts = doc_counts.clamp(min=1.0)
                mean_scores = doc_sums / doc_counts
                # TODO: 能不能不要重新产生一个local_doc_scores？
                local_doc_scores = torch.where(doc_counts > 0, mean_scores, local_doc_scores)

                del doc_sums, doc_counts, mean_scores
            else:
                raise ValueError(f"Invalid chunk reduction method: {self.chunk_reduce_method}")
            
            # 4. Global Scatter (Local Doc -> Global Doc)
            global_indices = slice_desc.original_doc_ids.expand(bsz, -1)
            global_doc_scores.scatter_reduce_(
                dim=1, index=global_indices, src=local_doc_scores, 
                reduce="amax", include_self=True
            )
            del attn_scores, max_scores_per_chunk, local_doc_scores

        # --- Phase 2: Top-K 选择 (使用 topk 替代 sort) ---
        
        # 优化点：topk 比 sort 快得多
        final_scores, batch_selected_doc_ids = torch.topk(global_doc_scores, k=min(self.model_config.doc_top_k, global_doc_scores.shape[1]), dim=1) # [bsz, k]
        batch_selected_global_doc_ids = self.block_desc.create_global_doc_ids(batch_selected_doc_ids)
        return final_scores, batch_selected_doc_ids, batch_selected_global_doc_ids

    def prefill_stage2_1(self, layer_idx: int, query_states: torch.Tensor, query_mask: torch.Tensor):
        # ... [Timer Start] ...
        
        # 优化点：Mask可以稍后处理，先让 Query 进计算
        
        bsz, num_heads, seq_len, head_dim = query_states.shape
        num_kv_groups = self.num_key_value_groups 
        num_kv_heads = num_heads // num_kv_groups
        min_val = torch.finfo(query_states.dtype).min
        routing_mask = ~query_mask
        q_lens = None
        
        # 1. View 重塑 (Zero-copy)
        # TODO: scaling 计算
        scaling = self.scaling
        query_states = query_states.view(bsz, num_kv_heads, num_kv_groups, seq_len, head_dim) * scaling

        # 初始化分数表
        # 严重注意：请确认 nr_docs 是否真的等于 max_doc_id + 1。
        # 如果 pool_ids 是全局唯一的（例如加了 bias 100000），这里的 total_docs 必须能容纳最大的 ID。
        total_docs = self.block_desc.nr_docs + 1 
        global_doc_scores = torch.full((bsz, total_docs), min_val, dtype=query_states.dtype, device=self.device)

        # --- Phase 1: 打分 (移除 Query Batch Loop) ---
        
        # 针对 A100，如果 bsz < 256，直接全量计算收益最高
        # 我们直接遍历 slice，不再对 query 进行切片
        
        for slice_desc, k_slice_t in zip(self.slice_desc, self.k_slices[layer_idx]):
            # k_slice_t: [1, nhead, 1, hdim, chunk] (GPU)
            
            # 1. Matmul (Compute Bound)
            # [bsz, nhead, groups, seqlen, hdim] @ [1, nhead, 1, hdim, chunk] 
            # -> [bsz, nhead, groups, seqlen, chunk]
            attn_scores = torch.matmul(query_states, k_slice_t)
            
            # 2. Mask & Reduce (Memory Bound)
            attn_scores.masked_fill_(routing_mask, min_val)
            # Flatten & Reduce heads/seqlen
            # flatten(1,2): [bsz, nhead*groups, seqlen, chunk]
            # max(dim=2)[0]: [bsz, nhead*groups, chunk]
            # max(dim=1)[0]: [bsz, chunk]
            # max_scores_per_chunk = attn_scores.flatten(1, 2).max(dim=2)[0].max(dim=1)[0] # [bsz, chunk]
            max_scores_per_chunk = attn_scores.flatten(1, 2)

            # [bsz, nhead*groups, chunk]
            if self.head_reduce_method == "max":
                max_scores_per_chunk = max_scores_per_chunk.max(dim=1).values
            elif self.head_reduce_method == "mean":
                max_scores_per_chunk = max_scores_per_chunk.mean(dim=1)
            else:
                raise NotImplementedError(f"Unsupported head reduce method: {self.head_reduce_method}")
            
            if self.query_reduce_method == "max":
                max_scores_per_chunk = max_scores_per_chunk.max(dim=1).values
            elif self.query_reduce_method == "mean":
                # 1. 调整 mask 形状以匹配 max_scores_per_chunk
                # [bsz, 1, seqlen, 1] -> [bsz, seqlen, 1]
                mask = query_mask.squeeze(1) 

                # 2. 在求和前，将无效位置（min_val）归零
                # 使用 torch.where 比乘法更安全，且能有效处理极小的 min_val
                masked_sum = torch.where(mask, max_scores_per_chunk, torch.zeros_like(max_scores_per_chunk)).sum(dim=1)

                # 3. 计算每个 sequence 中的有效 token 数量
                # 防止除以 0，使用 clamp 或加一个极小值 eps
                valid_counts = mask.sum(dim=1).clamp(min=1.0) 

                # 4. 得到最终均值 [bsz, chunk]
                max_scores_per_chunk = masked_sum / valid_counts
            elif self.query_reduce_method == "last":
                if q_lens is None:
                    q_lens = query_mask.squeeze(3).squeeze(1).sum(dim=1).long()
                
                # 边界保护: 如果 length 为 0 (异常情况), clamp 到 0
                last_indices = (q_lens - 1).clamp(min=0)
                
                # 3. 调整形状用于 gather
                # [B, 1, C]
                gather_idx = last_indices.view(bsz, 1, 1).expand(-1, 1, slice_desc.nr_chunks)
                
                # 4. 提取
                max_scores_per_chunk = max_scores_per_chunk.gather(1, gather_idx).squeeze(1)
            else:
                raise NotImplementedError(f"Unsupported query reduce method: {self.query_reduce_method}")

            # 3. Local Scatter (Chunk -> Doc)
            # slice_desc.local_doc_ids_0 已经在 GPU 上预计算好了
            # [1, chunk] -> [bsz, chunk]
            scatter_indices = slice_desc.local_doc_ids_0.expand(bsz, -1)
            
            local_doc_scores = torch.full((bsz, slice_desc.nr_docs), -float('inf'), device=self.device, dtype=query_states.dtype)
            if self.chunk_reduce_method == "max":
                local_doc_scores.scatter_reduce_(
                    dim=1, index=scatter_indices, src=max_scores_per_chunk, 
                    reduce="amax", include_self=True
                )
            
            elif self.chunk_reduce_method == "mean":
                doc_sums = torch.zeros_like(local_doc_scores)
                doc_sums = doc_sums.scatter_reduce(
                    dim=1, index=scatter_indices, src=max_scores_per_chunk, 
                    reduce="sum", include_self=False
                )
                doc_counts = torch.zeros_like(local_doc_scores)
                ones = torch.ones_like(max_scores_per_chunk)
                doc_counts = doc_counts.scatter_reduce(
                    dim=1, index=scatter_indices, src=ones, 
                    reduce="sum", include_self=False
                )
                doc_counts = doc_counts.clamp(min=1.0)
                mean_scores = doc_sums / doc_counts
                # TODO: 能不能不要重新产生一个local_doc_scores？
                local_doc_scores = torch.where(doc_counts > 0, mean_scores, local_doc_scores)

                del doc_sums, doc_counts, mean_scores
            else:
                raise ValueError(f"Invalid chunk reduction method: {self.chunk_reduce_method}")
            
            # 4. Global Scatter (Local Doc -> Global Doc)
            global_indices = slice_desc.original_doc_ids.expand(bsz, -1)
            global_doc_scores.scatter_reduce_(
                dim=1, index=global_indices, src=local_doc_scores, 
                reduce="amax", include_self=True
            )
            del attn_scores, max_scores_per_chunk, local_doc_scores

        # --- Phase 2: Top-K 选择 (使用 topk 替代 sort) ---
        
        # 优化点：topk 比 sort 快得多
        final_scores, batch_selected_doc_ids = torch.topk(global_doc_scores, k=self.model_config.doc_top_k, dim=1) # [bsz, k]
        
        required_doc_mask = torch.zeros(total_docs, dtype=torch.bool, device=self.device)
        required_doc_mask.scatter_(0, batch_selected_doc_ids.flatten(), True)
        
        if not required_doc_mask.any():
             return None, None, None, None

        # --- Phase 3: 异步并行提取 (Async Retrieval) ---
        
        block = self.blocks[layer_idx]
        pooled_doc_ids = self.block_desc.pool_ids # [num_chunks]
        
        # 计算 Chunk Mask (GPU)
        required_chunk_mask = required_doc_mask[pooled_doc_ids] # [num_chunks] Bool

        # 确保先处理 K
        k_selected = block.k[:, :, required_chunk_mask, :]
        pooled_doc_ids_selected = self.block_desc.create_global_doc_ids(pooled_doc_ids[required_chunk_mask])

        cpu_indices = torch.where(required_chunk_mask)[0].to("cpu")
        if rk_selected.is_cpu:
            rk_selected = block.rk[:, :, cpu_indices, :]
            rk_selected = rk_selected.to(self.device, non_blocking=True)
        else:
            rk_selected = block.rk[:, :, required_chunk_mask, :]

        v_selected = block.v[:, :, cpu_indices, :].to(self.device, non_blocking=True)

        # 由于 K 和 IDs 用了 non_blocking，我们需要在使用它们之前确保传输完成
        # 这里的 synchronize 确保 k_retrieved 和 pooled_doc_ids_retrieved 数据已就绪
        torch.cuda.current_stream().synchronize()

        return final_scores, k_selected, rk_selected, v_selected, pooled_doc_ids_selected

    def gpu_select(self, scores, k, rk, v, pooled_doc_ids, total_docs):
        final_scores, batch_selected_doc_ids = torch.topk(scores, k=self.model_config.doc_top_k, dim=1) # [bsz, k]
        
        required_doc_mask = torch.zeros(total_docs, dtype=torch.bool, device=self.device)
        required_doc_mask.scatter_(0, batch_selected_doc_ids.flatten(), True)
        
        # 计算 Chunk Mask (GPU)
        required_chunk_mask = required_doc_mask[pooled_doc_ids] # [num_chunks] Bool

        # 确保先处理 K
        k_selected = k[:, :, required_chunk_mask, :]
        pooled_doc_ids_selected = pooled_doc_ids[required_chunk_mask]
        rk_selected = rk[:, :, required_chunk_mask, :]
        v_selected = v[:, :, required_chunk_mask, :]

        return final_scores, k_selected, rk_selected, v_selected, pooled_doc_ids_selected


class MSAService(Memory, MemoryClientBase):
    def __init__(self,
                 gpu_id: int,
                 generate_config: GenerateConfig,
                 model_config: ModelConfig,
                 memory_config: MemoryConfig
                 ):
        self.world_size = generate_config.world

        # 初始化 NCCL
        # 注意：这里假设是在单机多卡环境下运行，使用 localhost
        if "MASTER_ADDR" not in os.environ:
            os.environ["MASTER_ADDR"] = "localhost"
        if "MASTER_PORT" not in os.environ:
            os.environ["MASTER_PORT"] = "29500"
        
        # 初始化进程组
        # timeout 设置稍长一点，防止初始化时某些进程慢导致超时
        dist.init_process_group(
            "nccl", 
            rank=gpu_id, 
            world_size=self.world_size
        )
        print(f"[GPU {gpu_id}] NCCL Initialized successfully.")

        super().__init__(gpu_id, generate_config, model_config, memory_config)

        self.generate_kwarg = self._create_generate_args()

    def _create_generate_args(self):
        generate_kwarg = {
            "do_sample": True,
            "top_p": self.generate_config.top_p,
            "temperature": self.generate_config.temperature,
            "max_new_tokens": self.generate_config.max_generate_tokens,
        }
        if self.generate_config.temperature == 0.0:
            generate_kwarg["do_sample"] = False
            generate_kwarg["temperature"] = None
            generate_kwarg["top_p"] = None
            generate_kwarg["top_k"] = None
        return generate_kwarg


    def setup_memory_client(self, model):
        def _setup_module(module):
            if hasattr(module, 'set_memory_client') and callable(module.set_memory_client):
                try:
                    module.set_memory_client(self)
                except Exception as e:
                    print(f"✗ 调用 {module.__class__.__name__} 的 set_memory_client() 失败: {e}")
        
        # 递归遍历所有模块
        for _, module in model.named_modules():
            _setup_module(module)

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

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

        self.model.eval()
        self.setup_memory_client(self.model)

    
    def _gather_querys(self, query_states: torch.Tensor, query_mask: torch.Tensor):
        """处理变长 seqlen：填充到统一长度并 all_gather"""
        bsz, nhead, seqlen, hdim = query_states.shape
        num_kv_groups = self.num_key_value_groups 
        num_kv_heads = nhead // num_kv_groups
        world_size = self.world_size
        device = self.device
        if self.scaling > 0:
            query_states = query_states * self.scaling

        # --- Step 1: 准备全量 Query (分布式 All-Gather) ---
        m_for_search = query_mask.view(bsz, 1, 1, seqlen, 1)
        q_for_search = query_states.view(bsz, num_kv_heads, num_kv_groups, seqlen, hdim)

        # --- Step 2: 处理变长 Seqlen (隐式同步处理) ---
        local_seqlen = torch.tensor([seqlen], device=device, dtype=torch.long)
        all_seqlens = [torch.empty(1, device=device, dtype=torch.long) for _ in range(world_size)]
        dist.all_gather(all_seqlens, local_seqlen)
        
        # 此处 s.item() 确实会触发同步，但这是为了在 CPU 端分配 all_queries 的 List
        seqlens = [s.item() for s in all_seqlens]
        max_seqlen = max(seqlens)

        # --- Step 3: 预分配连续内存进行 All-Gather (解决 Contiguous 问题) ---
        # 我们预先分配一个大的连续 Buffer，然后把 local 数据 copy 进去
        # 这样 gather 出来的数据天生就是连续的，不需要调 .contiguous()
        def get_padded_gather_list(tensor, target_seqlen):
            # 构造 padded 形状
            padded_shape = list(tensor.shape)
            padded_shape[3] = target_seqlen # 修改 seqlen 维度
            
            # 分配全量列表
            gather_list = [torch.zeros(padded_shape, dtype=tensor.dtype, device=device) for _ in range(world_size)]
            
            # 将本地数据 copy 到 gather_list 中对应的位置
            # 这里 slice 操作 + copy_ 比先 pad 再 gather 更节省临时显存
            gather_list[self.gpu_id][:, :, :, :seqlen, :].copy_(tensor)
            dist.all_gather(gather_list, gather_list[self.gpu_id])
            return gather_list

        all_queries = get_padded_gather_list(q_for_search, max_seqlen)
        all_masks = get_padded_gather_list(m_for_search, max_seqlen)

        return all_queries, all_masks, seqlens

    def doc_query(self, query_states: torch.Tensor, query_mask: torch.Tensor, layer_idx: int):
        """
        集成版：检索 + 跨卡提取 + 查表加速 + 结果归集
        返回: 
            final_k, final_v: [H, Total_Selected_C, D] 格式，直接用于 Flash Attn 散射
            final_scores: [B, TopK] 文档检索分数
            num_selected_chunks_per_sample: [B] 每个 Batch 选中的 chunk 总数
        """
        world_size = self.world_size
        bsz, nhead, seqlen, hdim = query_states.shape
        num_kv_groups = self.num_key_value_groups 
        num_kv_heads = nhead // num_kv_groups
        device = self.device
        dtype = query_states.dtype
        top_k = self.model_config.doc_top_k

        all_queries, all_masks, seqlens = self._gather_querys(query_states, query_mask)

        # --- Step 2: 循环处理各个 Rank 的请求并本地打分 ---
        scores_list, combined_ids_list = [], []
        for target_rank in range(world_size):
            seqlen = seqlens[target_rank]
            q_batch = all_queries[target_rank][:, :, :seqlen, :]
            m_batch = all_masks[target_rank][:, :, :seqlen, :]
            # q_batch 来自 target_rank，我们需要在当前 GPU 上寻找最佳匹配
            l_scores, l_ids, g_ids = self.prefill_stage2(layer_idx, q_batch, m_batch)
            scores_list.append(l_scores)     # [bsz, top_k]
            # 将 local_ids 和 global_ids 打包在最后一个维度，减少一次 all-to-all
            combined_ids_list.append(torch.stack([l_ids, g_ids], dim=-1))

        # --- Step 3: 元数据交换 (All-to-All) ---
        # 发送：我为各 Rank 算的 Top-K；接收：各 Rank 为“我”算的 Top-K
        send_scores = torch.stack(scores_list, dim=0) # [world_size, bsz, top_k]
        send_combined_ids = torch.stack(combined_ids_list, dim=0) # [world_size, bsz, top_k, 2]
        
        recv_scores = torch.empty_like(send_scores)
        recv_combined_ids = torch.empty_like(send_combined_ids)
        
        dist.all_to_all_single(recv_scores, send_scores)
        dist.all_to_all_single(recv_combined_ids, send_combined_ids)
        
        # --- Step 4: 全局 Top-K 决策 ---
        # shapes: [world_size, bsz, top_k] -> [bsz, world_size * top_k]
        candidate_scores = recv_scores.transpose(0, 1).reshape(bsz, -1)
        candidate_local_ids = recv_combined_ids[..., 0].transpose(0, 1).reshape(bsz, -1)
        candidate_global_ids = recv_combined_ids[..., 1].transpose(0, 1).reshape(bsz, -1)
        
        origin_ranks = torch.arange(world_size, device=device).view(world_size, 1, 1).expand(-1, bsz, top_k) # [world_size, bsz, top_k]
        origin_ranks = origin_ranks.transpose(0, 1).reshape(bsz, -1) # [bsz, world_size * top_k]

        # [bsz, top_k]
        final_scores, final_indices = torch.topk(candidate_scores, k=top_k, dim=1)
        winner_local_ids = candidate_local_ids.gather(1, final_indices)
        winner_global_ids = candidate_global_ids.gather(1, final_indices)
        winner_ranks = origin_ranks.gather(1, final_indices)

        # 屏蔽掉无效分数
        valid_mask = final_scores > -1e9
        winner_local_ids = winner_local_ids.masked_fill(~valid_mask, -1)

        # --- Step 5: 任务下发 (Request) ---
        # 告诉其他 Rank：我最终选了你哪几个 local_id
        task_mask = torch.arange(world_size, device=device).view(world_size, 1, 1)
        tasks_to_send = torch.where(winner_ranks == task_mask, winner_local_ids, -1)
        remote_tasks = torch.empty_like(tasks_to_send)
        dist.all_to_all_single(remote_tasks, tasks_to_send)
        remote_tasks_cpu = remote_tasks.to("cpu")
        
        # --- Step 6: 提取 KV (核心：针对百万级数据的查表加速) ---
        
        block = self.blocks[layer_idx]
        desc = self.block_desc

        # 1. 预处理：收集每个 Rank 的请求元数据 (在 CPU 上完成，极快)
        all_rank_local_ids = []
        chunks_per_rank = []  # 记录每个 Rank 分得的 chunk 总数，用于后续拆分
        all_rank_u_lens = []  # 记录每个 Rank 选中的每个 doc 的长度

        for r in range(world_size):
            requested_ids = remote_tasks_cpu[r]
            # 获取该 Rank 选中的唯一 local_doc_ids
            unique_req_ids = requested_ids[requested_ids != -1].unique()
            all_rank_local_ids.append(unique_req_ids)
            
            if unique_req_ids.numel() > 0:
                u_lens = desc.doc_lens_cpu[unique_req_ids]
                all_rank_u_lens.append(u_lens)
                chunks_per_rank.append(u_lens.sum().item())
            else:
                all_rank_u_lens.append(torch.empty(0, dtype=torch.long))
                chunks_per_rank.append(0)

        # 2. 构造全局批次索引 (One-shot Index Generation)
        flat_unique_ids = torch.cat(all_rank_local_ids)
        total_chunks_all_ranks = sum(chunks_per_rank)

        # 直接在 fetch 过程中完成 K 和 V 的合并
        # 预分配用于 all-to-all 的连续显存
        # total_chunks 是当前 Rank 需要发给所有其他 Rank 的 chunk 总和
        send_kv_combined = torch.empty((total_chunks_all_ranks, 2, num_kv_heads, hdim), device=device, dtype=dtype)

        if total_chunks_all_ranks > 0:
            # 批量获取所有选定文档的偏移和长度
            batch_u_offsets = desc.doc_offsets_cpu[flat_unique_ids]
            batch_u_lens = torch.cat(all_rank_u_lens)
            
            # 向量化生成全局索引列表
            base_indices = torch.repeat_interleave(batch_u_offsets, batch_u_lens)
            inner_seq = torch.arange(total_chunks_all_ranks) - torch.repeat_interleave(
                torch.cumsum(batch_u_lens, dim=0) - batch_u_lens, batch_u_lens
            )
            all_cpu_indices = base_indices + inner_seq
            all_gpu_indices = all_cpu_indices.to(device, non_blocking=True)
            
            # 3. 一次性提取并搬运 (核心性能提升点)
            # 由于 block.k/v 是 pinned，这一步是满速 DMA 异步传输
            # [1, H, C_all, D]
            def fetch(tensor, slot_idx):
                # tensor: [1, H, C, D] -> index_select -> [1, H, Sel_C, D]
                if tensor.is_cpu:
                    selected = tensor.index_select(2, all_cpu_indices).to(device, non_blocking=True)
                else:
                    selected = tensor.index_select(2, all_gpu_indices)
                send_kv_combined[:, slot_idx, :, :].copy_(selected.squeeze(0).transpose(0, 1), non_blocking=True)

            fetch(block.k, 0)
            fetch(block.v, 1)
            
            # 批量处理全局 Doc ID 映射
            all_g_ids_rep = desc.doc_ids_cpu[flat_unique_ids].repeat_interleave(batch_u_lens).to(device, non_blocking=True)
        else:
            # 全量空处理
            all_g_ids_rep = torch.empty(0, device=device, dtype=torch.long)

        send_num_chunks = torch.tensor(chunks_per_rank, device=device)

        # --- Step 7: 变长数据交换 (All-to-All) ---
        recv_num_chunks = torch.zeros_like(send_num_chunks)
        dist.all_to_all_single(recv_num_chunks, send_num_chunks)
        torch.cuda.current_stream().synchronize() # 确保 fetch(CPU) 已完成同步
        
        total_recv = recv_num_chunks.sum().item()
        recv_kv_flat = torch.empty(total_recv, 2, num_kv_heads, hdim, dtype=dtype, device=device)
        recv_g_ids_flat = torch.empty(total_recv, dtype=torch.long, device=device)
        
        in_splits = send_num_chunks.tolist()
        out_splits = recv_num_chunks.tolist()

        dist.all_to_all_single(recv_kv_flat, send_kv_combined, out_splits, in_splits)
        dist.all_to_all_single(recv_g_ids_flat, all_g_ids_rep, out_splits, in_splits)
        torch.cuda.current_stream().synchronize() 

        # --- Step 8: 后期处理集成 (按 Batch 重排归集) ---
        # 构造匹配矩阵 [B, TopK, Total_Recv_Chunks]
        # 利用广播找到每个 Batch Item 选中的 chunk
        # TODO: 这里可能会爆显存，后续考虑分块处理
        match_matrix = (winner_global_ids.unsqueeze(-1) == recv_g_ids_flat.unsqueeze(0).unsqueeze(0))
        mask_per_batch = match_matrix.any(dim=1) # [B, Total_Recv_Chunks]

        # 假设 recv_k_flat 形状为 [C, 1, H, D]
        # 1. 调整维度，使 B 和 C 排在最前面，方便 mask 索引
        # unsqueeze(0) -> [1, C, 1, H, D]
        # expand(bsz, ...) -> [bsz, C, 1, H, D]
        candidate_k_exp = recv_kv_flat[:,0:1, :, :].unsqueeze(0).expand(bsz, -1, -1, -1, -1) 
        candidate_v_exp = recv_kv_flat[:,1:2, :, :].unsqueeze(0).expand(bsz, -1, -1, -1, -1) 

        # 2. 利用布尔索引提取。
        # candidate_k_exp[mask_per_batch] 会根据 [B, C] 的 True 位置提取对应的 [1, H, D]
        # 结果形状: [Total_Selected_In_Batch, 1, H, D]
        final_k_raw = candidate_k_exp[mask_per_batch]
        final_v_raw = candidate_v_exp[mask_per_batch]

        # 3. 统计数量（这部分没问题）
        num_selected_chunks_per_sample = mask_per_batch.sum(dim=1)

        # 4. 适配输出格式 [H, Total_Selected_C, D]
        # 先 squeeze(1) 去掉那个大小为 1 的维度 -> [Total_Selected, H, D]
        # 再 permute(1, 0, 2) 交换 Total_Selected 和 H -> [H, Total_Selected, D]
        final_k_to_scatter = final_k_raw.squeeze(1).permute(1, 0, 2).contiguous()
        final_v_to_scatter = final_v_raw.squeeze(1).permute(1, 0, 2).contiguous()

        return final_k_to_scatter, final_v_to_scatter, final_scores, num_selected_chunks_per_sample, winner_global_ids

    def generate(self, req: GenerateRequest) -> GenerateResponse:
        # 转换为tensor
        input_ids_tensor = torch.LongTensor(req.input_ids).to(self.device)
        attention_mask_tensor = torch.LongTensor(req.attention_mask).to(self.device)
        doc_ids_tensor = torch.LongTensor(req.doc_ids).to(self.device)
        position_ids_tensor = torch.LongTensor(req.positions).to(self.device)

        past_key_values: CustomDynamicCache = create_cache()
        past_key_values.meta["require_recall_topk"] = req.require_recall_topk
        past_key_values.meta["qa_mode"] = self.generate_config.qa_mode
        past_key_values.meta["max_generate_tokens"] = self.generate_config.max_generate_tokens
        if self.generate_config.qa_mode:
            past_key_values.meta["tokenizer"] = self.tokenizer
            past_key_values.meta["idx_to_doc"] = self.idx_to_doc
            past_key_values.meta["pattern"] = r"\[(\d+)\]"
            past_key_values.meta["response_string"] = ['' for _ in range(len(req.input_ids))]
            

        # 准备generate kwargs
        for layer_idx in range(self.msa_model_config.num_hidden_layers):
            past_key_values.record_kwargs(layer_idx, {"stage": "prefill_stage2"})

        generate_kwargs = self.generate_kwarg.copy()
        generate_kwargs["past_key_values"] = past_key_values

        inputs = {
            "input_ids": input_ids_tensor,
            "attention_mask": attention_mask_tensor,
            "doc_ids": doc_ids_tensor,
            "use_cache": True,
            "position_ids": position_ids_tensor,
        }

        with torch.no_grad():
            generated_ids = self.model.generate(
                **inputs,
                **generate_kwargs,
            )
        generated_seqs = self.tokenizer.batch_decode(generated_ids, skip_special_token=True)
        if req.require_recall_topk:
            recall_topks = {layer: args['recall_topk'] for layer, args in past_key_values.cache_kwargs.items() if 'recall_topk' in args}
        else:
            recall_topks = None

        return GenerateResponse(msg_id=req.msg_id,
                                seq_id=req.seq_id,
                                gpu_id=self.gpu_id,
                                generated_texts=generated_seqs, 
                                recall_topk=recall_topks)



class MSAEngine:
    def __init__(self,
                 generate_config: GenerateConfig,
                 model_config: ModelConfig,
                 memory_config: MemoryConfig
                 ):
                 
        # NCCL require global cuda devices, so wen change the device setting here
        # user should set CUDA_VISIBLE_DEVICES envionment vars if they wish to select specific devices
        generate_config.devices = list(range(torch.cuda.device_count()))

        self.device_ids = generate_config.devices
        self.world_size = generate_config.world
        self.generate_config = generate_config
        self.model_config = model_config
        self.memory_config = memory_config

        self.worker_processes = []
        self.worker_qs: Dict[int, mp.Queue] = {}
        self.response_queue = mp.Queue()
        self.request_queue = mp.Queue()
        self.start_workers()

        self.msg_id = 0
        self.nr_active_req = 0
        self.limiter = RequestLimiter(4)
        self.requests: Dict[int, GenerateStub] = {} # msgid -> stub
        self.lock = threading.Lock()
        self.sync_event = threading.Event()
        self.sync_rsp = None

        self.initialize()
    
    def initialize(self):
        """初始化Memory服务"""
        print("Initializing Memory service...")

        # 加载tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_config.model_path)
        self._prepare_template()

        # 加载memory文件
        self._load_memory_file()

        # 先等待所有的worker都就绪了再发送 blocks
        self._wait_ready_signal()
        self._process_buckets()

        # 必须要在bucket处理完成后才启动线程
        self.running = True
        self.recv_rsp_thread = threading.Thread(target=self.receive_response, args=(), daemon=True)
        self.recv_rsp_thread.start()

        self.recv_req_thread = threading.Thread(target=self.receive_request, args=(), daemon=True)
        self.recv_req_thread.start()

        print("MSAEngine initialized")

        
    def _worker_all_gather(self, cmdname: str, timeout=600):
        collected_count = 0
        responses = []

        while collected_count < len(self.worker_qs):
            try:
                response = self.response_queue.get(timeout=timeout)
                assert response.name == cmdname, f"worker all gather got response {response.name} expect {cmdname}"
                responses.append(response)
                collected_count += 1

            except queue.Empty:
                continue

        responses.sort(key=lambda x: x.gpu_id)
        return responses

    def _wait_ready_signal(self):
        collected_count = 0

        while collected_count < len(self.worker_qs):
            ProtocolConstants.expect(self.response_queue, MEMORY_WORKER_READY)
            collected_count += 1

    @staticmethod
    def balanced_bucket_partition(docs: List[Document], n_buckets: int) -> List[List[Document]]:
        """
        全局均衡划分算法：先全局分blocks，再分配到buckets,
        这里并不需要每个 block 必须是最多max_chunk_per_block个 chunk，
        而是根据max_chunk_per_block这个数大致决定分成多少 block（GPU 个数的倍数），然后将这些
        block 做得尽可能 chunk 均衡，最后分配给各个 GPU

        Returns:
                List[List[Document]]: 每个 bucket 分配到的文档列表
        """

        # 全局 blocks 数量
        bucket_docs: List[List[Document]] = [[] for _ in range(n_buckets)] # 每个 bucket 所含的 doc 列表
        bucket_chunk_count = [0] * n_buckets

        def next_bucket():
            return bucket_chunk_count.index(min(bucket_chunk_count))

        for doc in docs:
            bucket_idx = next_bucket()
            bucket_docs[bucket_idx].append(doc)
            bucket_chunk_count[bucket_idx] += doc.num_chunks


        return bucket_docs

    def _sort_reference(self, docs: List[str]) -> Tuple[List[str], List[int], List[int], List[List[int]]]:
        """
        对文档进行重排序并且分 bucket和block，分配到一个 GPU 上的文档被称为 bucket，
        bucket 和 bucket 之间的 chunk 数量尽可能均衡
        Args:
            docs: 文档数据列表

        Returns:
                List[List[Document]]: 每个 bucket 分配到的文档列表
        """
        

        documents: List[Document] = []
        kernel_sz = self.model_config.pooling_kernel_size
        for idx, doc in enumerate(docs):  # idx 是 block 内部的 doc 索引
            new_doc, doc_inputs = compose_input(doc, idx, self.tokenizer)
            length = len(doc_inputs["input_ids"])
            num_chunks = (length + kernel_sz - 1) // kernel_sz
            # print(f"doc {idx} str {len(doc)} id length: {length}, num_chunks: {num_chunks}")

            documents.append(Document(doc=doc, doc_id=idx, num_chunks=num_chunks))
        
        return MSAEngine.balanced_bucket_partition(documents, self.generate_config.world)

    def _load_memory_file(self):
        """加载memory文件"""
        print(f"Loading memory file: {self.memory_config.memory_file_path}")

        if self.memory_config.memory_file_path.endswith('.json'):
            with open(self.memory_config.memory_file_path, 'r') as f:
                data = json.load(f)
            # 处理JSON格式数据
            # 这里需要根据实际的JSON格式来实现
            context_list = data

        elif self.memory_config.memory_file_path.endswith('.pkl'):
            with open(self.memory_config.memory_file_path, 'rb') as f:
                reference_metas = pickle.load(f)

            context_list = list(reference_metas)

        else:
            raise ValueError(f"Unsupported file format: {self.memory_config.memory_file_path}")
        
        # debug function:
        # when env DEBUG_SCALE_MEMORY is set by MemoryFilePath:ScaleFactor
        # we will read from MemoryFilePath and scale it to (ScaleFactor+1) times of memory documents
        # if it is set by ScaleFactor, without MemoryFilePath, we will use self.memory_config.memory_file_path instead
        scale_path = os.environ.get("DEBUG_SCALE_MEMORY", "")
        if scale_path:
            g = scale_path.split(":")
            scale = 0
            if len(g) == 2:
                scale_path = g[0]
                scale=int(g[1])
            elif len(g) == 1:
                scale_path = self.memory_config.memory_file_path
                scale=int(g[0])
            else:
                print("invalid scaling, ignore: ", scale_path)
                
            if scale > 0:
                print(f">>>>>>>>>>>>>>>>>>>> DEBUG: scale from query file {scale_path} with scale {scale}", )
                from src.utils.scale import scale_memory
                scale_memory(context_list, scale_path, scale=scale)

        self.buckets: List[List[Document]] = self._sort_reference(context_list)
        self.docs: List[Document] = []
        for bucket in self.buckets:
            self.docs.extend(bucket)
        chunks = sum(doc.num_chunks for doc in self.docs)

        print(f"Loaded {len(self.docs)} memory document, total chunks: {chunks}")
        for idx, bucket in enumerate(self.buckets):
            print(f"bucket {idx} contains {len(bucket)} documents, {sum(doc.num_chunks for doc in bucket)} chunks")

    def _process_buckets(self):
        """
        发送 blocks 到各 worker
        """

        idx_to_doc = self.get_idx_to_doc()
        for idx, gpu_id in enumerate(self.generate_config.devices):
            ProtocolConstants.send(self.worker_qs[gpu_id],
                                   MEMORY_WORKER_BLOCKS,
                                   data=self.buckets[idx],
                                   block=False)
            ProtocolConstants.send(self.worker_qs[gpu_id],
                                   MEMORY_WORKER_IDX_TO_DOC,
                                   data=idx_to_doc,
                                   block=False)

        # 拿到所有的 doc id bias然后累积更新各 worker
        rsps: List[ReportDocID] = self._worker_all_gather( "report_doc_id", timeout=None)       

        # pooled_doc_id_bias = 0
        # for rsp in rsps:
        #     self.worker_qs[rsp.gpu_id].put(SetDocIDCmd(pooled_doc_id_bias=pooled_doc_id_bias))
        #     pooled_doc_id_bias += rsp.pooled_doc_id_bias
        # _ = self._worker_all_gather("set_doc_id")

    def receive_request(self):
        while self.running:
            try:
                gpu_id, res = self.request_queue.get(timeout=3)
                print(f"Received response from GPU {gpu_id}")
            except queue.Empty:
                continue

    def get_idx_to_doc(self):
        return {doc.doc_id: doc.doc for doc in self.docs}

    @staticmethod
    def service_main(gpu_id: int,
                     request_queue: mp.Queue,
                     response_queue: mp.Queue,
                     generate_config: GenerateConfig,
                     model_config: ModelConfig,
                     memory_config: MemoryConfig
                     ):
        # 1. 初始化 Service (包含 NCCL Init)
        service = MSAService(gpu_id, generate_config, model_config, memory_config)

        print(f"MSAService-{gpu_id} is ready")

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

        # wait parent to give me blocks to process
        docs: List[Document] = ProtocolConstants.expect(request_queue, MEMORY_WORKER_BLOCKS)
        idx_to_doc: Dict[int, str] = ProtocolConstants.expect(request_queue, MEMORY_WORKER_IDX_TO_DOC)
        service.save_idx_to_doc(idx_to_doc)

        service.generate_blocks(docs)
        del docs
        torch.cuda.empty_cache()

        service.load_model()

        # tell Memory my doc id bias
        response_queue.put(ReportDocID(gpu_id=gpu_id, pooled_doc_id_bias=service.get_max_local_pool_doc_id()), block=False)

        
        while True:
            try:

                cmd: CmdBase = request_queue.get()
                # 检查是否是终止信号
                if cmd is None:
                    if dist.is_initialized():
                        dist.destroy_process_group()
                    return

                elif cmd.name == "generate_request":
                    rsp = service.generate(cmd)
                    response_queue.put(rsp)
                    # torch.cuda.empty_cache() # make perf worse

                else:
                    print(f"GPU_ID {gpu_id} recv unknown cmd {cmd.name}")

            except Exception as e:
                print(f"[subprocess {gpu_id}] error: {e}")
                import traceback
                traceback.print_exc()


    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stop_workers()
        return False  # 不抑制异常


    def start_workers(self):
        # 设置启动方式，CUDA 必须用 spawn
        mp.set_start_method('spawn', force=True)

        for gpu_id in self.generate_config.devices:
            request_queue = mp.Queue()
            process = mp.Process(
                target=MSAEngine.service_main,
                args=(gpu_id, request_queue, self.response_queue, self.generate_config, self.model_config, self.memory_config),
                name=f"MSAService-{gpu_id}"
            )
            process.start()
            self.worker_processes.append(process)
            self.worker_qs[gpu_id] = request_queue

    def stop_workers(self):
        print("MSAEngine stop all services")
        self.running = False

        for q in self.worker_qs.values():
            q.put(None)
        
        self.worker_qs = {}

        # 等待所有worker退出
        for process in self.worker_processes:
            process.join()
        self.worker_processes = []

        self.recv_rsp_thread.join()
        self.recv_rsp_thread = None

        self.recv_req_thread.join()
        self.recv_req_thread = None

    def _validate_inputs(self, input_ids: List[List[int]]):
        if any(len(sub) == 0 for sub in input_ids):
            raise ValueError(f"empty query is not acceptable")

        if self.generate_config.max_seq_len > 0:
            seqlen = sum(sum(sub) for sub in input_ids)
            if seqlen > self.generate_config.max_seq_len:
                raise ValueError(f"total sequence length {seqlen} exceeds limitation {self.generate_config.max_seq_len}")

        if self.generate_config.max_query_seq_len > 0:
            max_seq_len = max(sum(sub) for sub in input_ids)
            if max_seq_len > self.generate_config.max_query_seq_len:
                raise ValueError(f"max query sequence length {max_seq_len} exceeds limitation {self.generate_config.max_query_seq_len}")

    def _prepare_template(self):
        response_head = "<|im_start|>"
        response_head_inputs = self.tokenizer(response_head, add_special_tokens=False)
        response_head_input_ids = response_head_inputs["input_ids"]
        response_head_attention_mask = response_head_inputs["attention_mask"]

        pad_token = self.tokenizer.pad_token
        pad_token_id = self.tokenizer.pad_token_id
        prompt_template = self.generate_config.template["prompt"].replace("{prompt}", pad_token)
        prompt_template_inputs = self.tokenizer(prompt_template, add_special_tokens=False)
        prompt_template_input_ids = prompt_template_inputs["input_ids"]
        prompt_template_attention_mask = prompt_template_inputs["attention_mask"]
        pad_index = prompt_template_input_ids.index(pad_token_id)
        template_tail_input_ids = prompt_template_input_ids[pad_index+1:]
        template_tail_attention_mask = prompt_template_attention_mask[pad_index+1:]

        self.prompt_tail_input_ids =  template_tail_input_ids + response_head_input_ids
        self.prompt_tail_attention_mask =  template_tail_attention_mask + response_head_attention_mask
        self.tail_doc_ids = [-2] * (len(template_tail_input_ids)) + [-1] * len(response_head_input_ids)
        
    def _apply_template(self, prompt):
        # question = "\n请根据以上历史文档信息，回答问题\n\n" + prompt + "\n" + f"请返回与问题有关的所有文档\n"
        # question = "\n请根据以上历史文档信息，回答问题\n\n" + prompt + "\n" + f"请返回与问题有关的1个文档\n"
        question = "\nPlease answer the question based on the above historical document information\n\n" + prompt + "\n" + f"Please return all documents related to the question\n"
        prompt_inputs = self.tokenizer(question, add_special_tokens=False)
        prompt_inputs["input_ids"], prompt_inputs["attention_mask"]

        input_ids = prompt_inputs["input_ids"] + self.prompt_tail_input_ids
        attention_mask = prompt_inputs["attention_mask"] + self.prompt_tail_attention_mask
        doc_ids = [0] * (len(prompt_inputs["input_ids"])) + self.tail_doc_ids
        return input_ids, attention_mask, doc_ids
    
    def _apply_template_regenerate(self, prompt):
        question = prompt.split("<regenerate>")[1]

        prompt_inputs = self.tokenizer(question, add_special_tokens=False)
        prompt_inputs["input_ids"], prompt_inputs["attention_mask"]

        input_ids = prompt_inputs["input_ids"]
        attention_mask = prompt_inputs["attention_mask"]

        # Build doc_ids: 0 for specific regions, -1 for the rest
        doc_ids = [-1] * len(prompt_inputs["input_ids"])

        # Region 1: tokens between the last "]<|object_ref_end|>[" and the first "<|object_ref_end|>" after it
        last_marker_start = question.rfind("]<|object_ref_end|>[")
        if last_marker_start != -1:
            search_start = last_marker_start + len("]<|object_ref_end|>[") - 1
            next_marker = question.find("<|object_ref_end|>", search_start)
            if next_marker != -1:
                region1_text = question[search_start:next_marker]
                region1_tokens = self.tokenizer(region1_text, add_special_tokens=False)["input_ids"]
                prefix_before_region1 = question[:search_start]
                prefix_tokens_len = len(self.tokenizer(prefix_before_region1, add_special_tokens=False)["input_ids"])
                for i in range(prefix_tokens_len, prefix_tokens_len + len(region1_tokens)):
                    if i < len(doc_ids):
                        doc_ids[i] = 0

        # Region 2: "Please return all documents related to the question" and everything before it
        anchor = "Please return all documents related to the question"
        anchor_pos = question.find(anchor)
        if anchor_pos != -1:
            prefix_with_anchor = question[:anchor_pos + len(anchor)]
            prefix_token_len = len(self.tokenizer(prefix_with_anchor, add_special_tokens=False)["input_ids"])
            for i in range(prefix_token_len):
                if i < len(doc_ids):
                    doc_ids[i] = 0

        return input_ids, attention_mask, doc_ids

    def receive_response(self):
        while self.running:
            try:
                rsp: GenerateResponse = self.response_queue.get(timeout=3)

                final_stub = None
                with self.limiter.lock:
                    assert rsp.msg_id in self.requests
                    stub = self.requests[rsp.msg_id]
                    stub.responses.append(rsp)
                    # print(f"Request {rsp.msg_id} got {len(stub.responses)} responses")
                    if len(stub.responses) == self.world_size:
                        final_stub = self.requests.pop(rsp.msg_id)
                
                if final_stub:
                    # print("Release")
                    self.limiter.release()
                    final_stub.respond()
                        
            except queue.Empty:
                continue

    def default_callback(self, texts: List[str], recall_topk, userdata):
        self.sync_rsp = (texts, recall_topk, userdata)
        self.sync_event.set()

    def generate(self,
                 prompts: Union[str, List[str]],
                 userdata=None,
                 require_recall_topk=False,
                 callback: GenerateReceiver=None):

        if isinstance(prompts, str):
            prompts = [prompts]
        
        world = self.generate_config.world
        bsz = (len(prompts) + world - 1) // world
        if self.generate_config.max_batch_size > 0 and bsz > self.generate_config.max_batch_size:
            raise ValueError(f"Batch size exceeds maximum allowed batch size {self.generate_config.max_batch_size}")
        # Important:
        #     如果输入的 prompts 不够 world size，则需要 pad，我们的方法是用最后一个prompt 填充
        pads = bsz * world - len(prompts)
        if pads > 0:
            prompts += [prompts[-1]] * pads

        batch_input_ids = []
        batch_attention_mask = []
        batch_doc_ids = []
        batch_positions = []

        # template + pooled docs
        current_position = 3 + self.model_config.doc_top_k

        for prompt in prompts:
            if "<regenerate>" in prompt:
                input_ids, attention_mask, doc_ids = self._apply_template_regenerate(prompt)
            else:
                input_ids, attention_mask, doc_ids = self._apply_template(prompt)
            batch_input_ids.append(input_ids)
            batch_attention_mask.append(attention_mask)
            batch_doc_ids.append(doc_ids)

            # 计算position ids
            position_ids = [current_position + i for i in range(len(input_ids))]
            batch_positions.append(position_ids)
            # current_position += len(input_ids)   # 为什么要累计长度？
        
        # Pad sequences to the same length
        padded_input_ids = []
        padded_attention_mask = []
        padded_doc_ids = []
        padded_position_ids = []

        pad_token_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id else self.tokenizer.eos_token_id
        # 使用 doc_end_id 作为 doc_ids 的 padding 值
        # doc_end_id = self.tokenizer("<|im_end|>", add_special_tokens=False)["input_ids"][0]

        for i in range(len(batch_input_ids)):
            if i % bsz == 0:
                max_length = max(len(ids) for ids in batch_input_ids[i:i+bsz])
            input_ids = batch_input_ids[i]
            attention_mask = batch_attention_mask[i]
            doc_ids_seq = batch_doc_ids[i]
            position_ids = batch_positions[i]

            # Pad to max_length
            pad_length = max_length - len(input_ids)
            # 改成left pad，方便position_ids迭代，可直接与自定义的generate兼容
            if pad_length > 0:
                input_ids = [pad_token_id] * pad_length + input_ids
                attention_mask = [0] * pad_length + attention_mask
                doc_ids_seq = [0] * pad_length + doc_ids_seq
                position_ids = [position_ids[0]] * pad_length + position_ids

            padded_input_ids.append(input_ids)
            padded_attention_mask.append(attention_mask)
            padded_doc_ids.append(doc_ids_seq)
            padded_position_ids.append(position_ids)

        # 只检查每个 batch 的第一个 query 就好
        self._validate_inputs(padded_input_ids[0:-1:bsz])


        # TODO: when callback is None
        self.limiter.acquire()

        sync = callback is None
        if callback is None:
            callback = self.default_callback

        self.msg_id += 1
        with self.limiter.lock:
            self.requests[self.msg_id] = GenerateStub(msg_id=self.msg_id,
                                                      userdata=userdata,
                                                      nr_dummy=pads,
                                                      responses=[],
                                                      callback=callback)

        for idx, gpu_id in enumerate(self.generate_config.devices):
            start, end = idx * bsz, (idx+1) * bsz
            req = GenerateRequest(msg_id=self.msg_id,
                                  seq_id=idx,
                                  prompts=prompts[start:end],
                                  input_ids=padded_input_ids[start:end],
                                  attention_mask=padded_attention_mask[start:end],
                                  doc_ids=padded_doc_ids[start:end],
                                  positions=padded_position_ids[start:end],
                                  require_recall_topk=require_recall_topk)
            # print(f"send GPU{gpu_id} batch {len(req.input_ids)}")
            self.worker_qs[gpu_id].put(req, block=False)

        if sync:
            self.sync_event.wait()
            self.sync_event.clear()
            return self.sync_rsp