Anonymous code release: MSA inference and evaluation

QUICK_START.md

@@ -0,0 +1,115 @@
+## Project Structure
+
+```
+MSA/
+├── scripts/
+│   ├── run_benchmarks.sh          # Run inference on benchmarks
+│   ├── calculate_llm_score.sh     # LLM-based answer evaluation
+│   └── resave_model.sh            # Convert base model to MSA format
+└── src/
+    ├── msa/                       # Core MSA implementation
+    │   ├── configuration_msa.py   # MSA model configuration
+    │   ├── memory_sparse_attention.py  # MemorySparseAttention layer
+    │   ├── model.py               # MSAForCausalLM / MSAModel
+    │   └── generate.py            # Generation logic
+    ├── config/
+    │   └── memory_config.py       # GenerateConfig, ModelConfig, MemoryConfig
+    ├── evaluation/
+    │   └── llm_judge.py           # LLM-based evaluation metrics
+    ├── app/
+    │   └── benchmark.py           # Benchmark runner
+    ├── utils/                     # GPU workers, caching, templates, etc.
+    ├── msa_service.py             # Multi-GPU inference engine (MSAEngine)
+    ├── prefill.py                 # Stage 1 prefill worker
+    ├── benchmarks.py              # Benchmark registry & specs
+    └── types.py                   # Core type definitions
+```
+
+## Installation
+
+**1. Create conda environment**
+
+```bash
+conda create -n msa python=3.12 -y
+conda activate msa
+```
+
+**2. Install dependencies**
+
+```bash
+pip install -r requirements.txt
+```
+
+<details>
+<summary>requirements.txt</summary>
+
+```
+torch==2.6.0
+torchvision==0.21.0
+transformers==4.51.3          # exact version required
+accelerate==1.0.1
+liger_kernel==0.5.10
+huggingface_hub==0.31.4       # must stay <1.0 for transformers 4.51.3
+datasets==3.1.0
+lmdb==1.6.2
+tqdm==4.67.1
+numpy==1.26.4
+pillow==11.2.1
+packaging==25.0
+nvidia-ml-py==12.575.51       # provides the `pynvml` module
+openai==1.79.0
+```
+
+</details>
+
+**3. Install Flash Attention**
+
+Option A — build from source:
+
+```bash
+pip install flash-attn==2.7.4.post1 --no-build-isolation
+```
+
+Option B — install prebuilt wheel (CUDA 12, Python 3.12):
+
+```bash
+wget -P /tmp https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl
+pip install /tmp/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl
+rm /tmp/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl
+```
+
+## Download
+
+**1. Download model**
+
+```bash
+mkdir ckpt
+pip install -U huggingface_hub==0.31.4
+export HF_ENDPOINT=https://hf-mirror.com
+huggingface-cli download --resume-download Anoy123423123/MSA-4B --local-dir ckpt/MSA-4B
+```
+
+**2. Download benchmarks**
+
+Benchmark data is hosted on [Anoy123423123/MSA-RAG-BENCHMARKS](https://huggingface.co/datasets/Anoy123423123/MSA-RAG-BENCHMARKS) and will be automatically downloaded to `data/` on first run, based on the benchmarks specified in `scripts/run_benchmarks.sh`. No manual download is needed.
+
+## Quick Start
+
+**1. Run inference on benchmarks**
+
+```bash
+bash scripts/run_benchmarks.sh eval_benchmark
+```
+
+**2. Compute LLM-based scores**
+
+```bash
+bash scripts/calculate_llm_score.sh eval_benchmark
+```
+
+## Supported Benchmarks
+
+| Category | Benchmark |
+|---|---|
+| Multi-hop QA | `2wikimultihopqa`, `hotpotqa`, `musique` |
+| Single-hop QA | `nature_questions`, `triviaqa_06M`, `triviaqa_10M`, `msmarco_v1`, `dureader`, `ms_100M`, `hipporag_narrative`, `hipporag_popqa` |

README.md

@@ -0,0 +1,291 @@
+---
+license: mit
+language:
+- en
+tags:
+- msa
+---
+# MSA: Memory Sparse Attention
+
+*A scalable, end-to-end trainable latent-memory framework for 100M-token contexts*
+
+**Anonymous code release for peer review.**
+
+| | |
+|---|---|
+| Code (this repo) | `Anoy123423123/MSA-Code` |
+| Model weights | [`Anoy123423123/MSA-4B`](https://huggingface.co/Anoy123423123/MSA-4B) |
+| Benchmark data | [`Anoy123423123/MSA-RAG-BENCHMARKS`](https://huggingface.co/datasets/Anoy123423123/MSA-RAG-BENCHMARKS) |
+
+Jump to [**Quick Start**](#-quick-start) to reproduce the reported numbers.
+
+## 📝 Abstract
+
+Long-term memory is essential for general intelligence, yet the **full attention** bottleneck constrains most LLMs' **effective context length** to **128K–1M**. Existing attempts — hybrid linear attention, fixed-size state memory (e.g., RNNs), and external storage like **RAG/agents** — either suffer rapid precision decay and latency growth at extreme scales, lack end-to-end differentiability or dynamic memory maintenance, or require complex pipelines. We present **Memory Sparse Attention (MSA)**: an **end-to-end trainable, scalable sparse latent-state memory** framework. Core ideas include:
+
+- **Scalable sparse attention** + **document-wise RoPE** (parallel/global) achieving **near-linear complexity** in both training and inference;
+- **KV cache compression** with a **Memory Parallel** inference engine to deliver **100M token** throughput on **2×A800** GPUs;
+- **Memory Interleave** for multi-round, multi-hop reasoning across scattered memory segments.
+
+On long-context QA and NIAH (Needle-in-a-Haystack) benchmarks, **MSA** surpasses same-backbone RAG, best-of-breed RAG stacks, and leading long-context models. Across an unprecedented **16K→100M token** range, MSA shows **< 9%** degradation, suggesting a practical path to **decouple memory capacity from reasoning**.
+
+> **Scaling from 16K→100M tokens**: MSA fuses top-k selection with sparse attention to remain end-to-end differentiable while allowing document decoupling at inference. On MS MARCO, MSA sustains **<9%** degradation and exhibits strong extrapolation.
+
+## ✨ Key Contributions
+
+- **Memory-Sparse Attention (MSA)**: an **end-to-end trainable**, **scalable sparse attention** layer with **document-wise RoPE**, realizing **O(L)** complexity and **<9%** degradation from **16K→100M tokens**.
+- **KV cache compression + Memory Parallel**: tiered storage (GPU-resident routing keys, CPU content K/V), distributed scoring, and on-demand transfers to enable **100M-token** inference on **2×A800**.
+- **Memory Interleave**: adaptive alternating "generative retrieval → context expansion → generation," significantly boosting **multi-hop** reasoning across documents.
+- **Comprehensive evaluation**: MSA outperforms same-backbone RAG, best-of-breed RAG pipelines, and top long-context models on long-context QA and NIAH, showing superior **stability** and **accuracy** at scale.
+
+---
+
+## 🧩 Overall Design
+
+### Architecture
+
+MSA integrates **retrieval and generation** into a single differentiable loop. Document latent states (**K/V/Kᵣ**) are **chunk-mean pooled** for compression. A **router projector** computes relevance via cosine similarity (mean-pooled over heads, then token-wise max), selects **Top‑k documents**, then concatenates their **compressed K/V** with the query's local **K/V** for autoregressive decoding. Routing applies only to **upper layers**; lower layers keep **independent document processing** for hierarchical alignment.
+
+- **Parallel (document-wise) RoPE**: Each document resets positions from 0, preventing position drift between **train-short** and **infer-long**, enabling 64k training to extrapolate to 100M.
+- **Global RoPE (active context)**: The query's starting index is **offset by k** (Top‑k retrieved blocks), preserving causal ordering: *background → query → generation*.
+
+Implemented in [`src/msa/memory_sparse_attention.py`](./src/msa/memory_sparse_attention.py) and [`src/msa/model.py`](./src/msa/model.py).
+
+---
+
+### Inference Pipeline
+
+MSA uses a **three-stage** pipeline:
+
+1. **Global Memory Encoding (offline)**: forward over the corpus to cache chunk-pooled **(K̄, V̄, K̄ᵣ)**.
+2. **Online Routing & Context Assembly**: project query to **Qᵣ**, match with **K̄ᵣ** to pick **Top‑k**, then load only the selected **K̄/V̄** and concatenate with local context.
+3. **Sparse Generation**: autoregress over the **sparse context**.
+
+**Memory Parallel** shards **K̄ᵣ** across GPUs (query broadcast → local scoring → global reduce). Content **K̄/V̄** stays in host DRAM and is **asynchronously fetched** when selected—balancing **VRAM** and **throughput** for **100M-token** deployment.
+
+Implemented in [`src/msa_service.py`](./src/msa_service.py) (`MSAEngine`, memory-parallel bucketing) and [`src/prefill.py`](./src/prefill.py) (stage-1 encoding).
+
+---
+
+## 🚀 Results
+
+> **Setup**
+> **QA**: 9 datasets (MS MARCO v1, NQ, DuReader, TriviaQA(10M), NarrativeQA, PopQA, 2WikiMultiHopQA, HotpotQA, MuSiQue), memory banks **277K→10M tokens**, metric: **LLM judge (0–5)**.
+> **NIAH (RULER)**: 8 subtasks, **32K→1M tokens**, report average accuracy.
+> **Backbone**: Qwen3‑4B‑Instruct‑2507. Compare to same-backbone RAG and best-of-breed RAG stacks (KaLMv2 + large generators, optional reranker).
+
+### Table 1: MSA vs same-backbone RAG (Qwen3‑4B)
+
+**Summary**: Average **3.760**, improving over standard RAG (**+16.0%**), RAG+rerank (**+11.5%**), and HippoRAG2 (**+14.8%**) using their best@k; MSA leads on all but NarrativeQA within the same-backbone group.
+
+| Dataset | Tokens | Qwen3-4B R@1 | R@5 | R@10 | Qwen3-4B (RR) R@1 | R@5 | R@10 | HippoRAG2 R@1 | R@5 | R@10 | MSA (adaptive) |
+|---------|--------|--------------|------|-------|----------------------|------|-------|------------------|------|--------|----------------|
+| MS MARCO v1 | 7.34M | 2.893 | 3.011 | 3.005 | 2.934 | <u>3.032</u> | 3.017 | 2.676 | 3.005 | 3.019 | **4.141** |
+| Natural Questions | 1.47M | 3.452 | 3.374 | 3.297 | <u>3.494</u> | 3.408 | 3.385 | 3.338 | 3.389 | 3.374 | **3.545** |
+| DuReader | 277K | 3.726 | 3.579 | 3.594 | <u>3.848</u> | 3.618 | 3.607 | 2.941 | 3.485 | 3.415 | **4.155** |
+| TriviaQA (10M) | 10M | 4.133 | 4.414 | 4.273 | 4.313 | 4.375 | 4.391 | 4.188 | <u>4.430</u> | 4.367 | **4.621** |
+| NarrativeQA | 538K | 1.611 | 2.567 | 2.860 | **3.638** | 3.492 | <u>3.536</u> | 1.959 | 2.628 | 2.655 | 3.395 |
+| PopQA | 1.18M | 2.959 | 3.273 | 3.299 | <u>3.315</u> | 3.264 | 3.266 | 3.111 | 3.249 | 3.249 | **3.433** |
+| 2WikiMultiHopQA | 722K | 1.065 | 3.055 | 3.136 | 1.187 | 3.057 | 3.159 | 1.045 | 3.180 | <u>3.330</u> | **4.280** |
+| HotpotQA | 1.35M | 2.252 | 3.582 | 3.787 | 2.642 | 3.990 | <u>4.022</u> | 3.230 | 3.770 | 3.970 | **4.061** |
+| MuSiQue | 1.41M | 0.936 | 1.752 | 1.928 | 1.144 | 1.960 | 1.965 | 1.020 | 1.907 | <u>2.095</u> | **2.211** |
+| **Average** | — | 2.559 | 3.179 | 3.242 | 2.946 | 3.355 | <u>3.372</u> | 2.612 | 3.227 | 3.275 | **3.760** |
+
+*Table 1: Same-backbone RAG vs MSA (@1/@5/@10 vs MSA @adaptive)*
+
+---
+
+### Table 2: MSA vs best-of-breed RAG (large backbones)
+
+**Summary**: Against **KaLMv2+Qwen3‑235B** and **KaLMv2+Llama‑3.3‑70B** (w/ and w/o reranking), MSA achieves the best score on **4/9** datasets and an average **3.760**, with relative gains of **+7.2%**, **+5.0%**, **+10.7%**, and **+5.4%** over the strongest configurations respectively. Gaps on a few datasets (e.g., MuSiQue) are largely attributable to parameter-count and intrinsic reasoning capacity.
+
+| Dataset | KaLMv2 + Qwen3‑235B R@1 | R@5 | R@10 | Qwen3‑235B (RR) R@1 | R@5 | R@10 | KaLMv2 + Llama‑3.3 R@1 | R@5 | R@10 | Llama‑3.3 (RR) R@1 | R@5 | R@10 | MSA (adaptive) |
+|---------|---------------------------|------|-------|--------------------------|------|-------|------------------------------|------|--------|-------------------------|------|--------|----------------|
+| MS MARCO v1 | 2.846 | <u>3.028</u> | 3.027 | 2.886 | 3.020 | 2.995 | 2.649 | 2.904 | 2.919 | 2.881 | 2.955 | 2.952 | **4.141** |
+| Natural Questions | <u>3.711</u> | 3.670 | 3.694 | 3.621 | 3.610 | 3.645 | 3.675 | 3.674 | 3.662 | **3.756** | 3.665 | 3.647 | 3.545 |
+| DuReader | 4.044 | 3.991 | 3.978 | 3.973 | 3.932 | 3.891 | <u>4.051</u> | 3.846 | 3.742 | 3.967 | 3.776 | 3.780 | **4.155** |
+| TriviaQA (10M) | 4.367 | 4.656 | 4.578 | 4.492 | 4.320 | 4.555 | 4.273 | **4.740** | <u>4.719</u> | 4.547 | 4.703 | 4.695 | 4.621 |
+| NarrativeQA | 1.413 | 2.130 | 2.427 | 3.212 | **3.427** | 3.375 | 1.290 | 2.123 | 2.382 | 3.150 | 3.263 | 3.317 | <u>3.395</u> |
+| PopQA | 2.810 | 3.347 | <u>3.396</u> | 3.268 | 3.380 | 3.376 | 2.787 | 3.298 | 3.305 | 3.337 | 3.384 | 3.362 | **3.433** |
+| 2WikiMultiHopQA | 2.646 | 3.579 | 3.582 | 1.855 | 3.381 | <u>3.583</u> | 1.339 | 3.263 | 3.445 | 1.651 | 3.332 | 3.541 | **4.280** |
+| HotpotQA | 3.497 | 4.090 | **4.225** | 3.341 | 4.141 | 4.194 | 3.070 | 3.896 | 4.127 | 3.428 | 4.145 | <u>4.203</u> | 4.061 |
+| MuSiQue | 1.988 | 2.462 | **2.647** | 1.801 | 2.522 | 2.605 | 1.704 | 2.317 | 2.258 | 1.895 | 2.462 | <u>2.614</u> | 2.211 |
+| **Average** | 3.036 | 3.439 | 3.506 | 3.161 | 3.526 | <u>3.580</u> | 2.760 | 3.340 | 3.396 | 3.179 | 3.521 | 3.568 | **3.760** |
+
+*Table 2: SOTA RAG stacks (strong retriever + large generator + optional reranker) vs MSA*
+
+---
+
+### RULER NIAH stability (32K→1M)
+
+**Summary**: MSA maintains **94.84%** at **1M tokens**. The unmodified backbone collapses beyond **128K** (down to **24.69% @1M**). Hybrid linear-attention long-context models degrade noticeably at **≥128K/256K**. External-memory agents (e.g., RL‑MemoryAgent‑14B) remain stable but are weaker in **absolute accuracy** and show steeper decay than MSA.
+
+---
+
+## Implementation Notes
+
+- **Training**: 158.95B-token continuous pretraining with **auxiliary routing loss**, followed by two-stage SFT (**8k→64k** curriculum).
+- **Ablations** (paper Table 4): curriculum extension, Memory Interleave, continuous pretraining, and injecting original text all contribute substantially; removing them causes **5%–37%** drops depending on task.
+- **Training cost**: MSA's near-linear complexity shows up directly in wall-clock time — full attention grows quadratically with context length while MSA stays close to linear.
+
+---
+
+## 🚀 Quick Start
+
+> This section is self-contained. [**QUICK_START.md**](./QUICK_START.md) additionally documents a
+> prebuilt `flash-attn` wheel, which avoids a slow source build.
+
+### Requirements
+
+| | |
+|---|---|
+| OS | Linux, x86-64 |
+| Python | 3.12 (tested on 3.12.9) |
+| CUDA | 12.x |
+| GPUs | **8× 80GB** for the default benchmark suite; **2× 80GB (A800)** is enough for the `ms_100M` 100M-token setting via Memory Parallel |
+| Disk | ~9 GB for weights + ~0.6 GB for benchmark data |
+
+`transformers` **must** be exactly `4.51.3` — the attention implementation targets that API and will not
+load on 4.52+ or on 5.x. The other pins in `requirements.txt` are the versions everything was tested with.
+
+### 1. Get the code
+
+```bash
+pip install -U "huggingface_hub==0.31.4"
+# If you are behind a mirror, uncomment the next line:
+# export HF_ENDPOINT=https://hf-mirror.com
+
+huggingface-cli download Anoy123423123/MSA-Code --repo-type=model --local-dir MSA
+cd MSA
+```
+
+> `git clone https://huggingface.co/Anoy123423123/MSA-Code` works too.
+
+### 2. Install dependencies
+
+```bash
+conda create -n msa python=3.12 -y
+conda activate msa
+pip install -r requirements.txt
+```
+
+`flash-attn` is installed separately because it needs `--no-build-isolation`:
+
+```bash
+pip install flash-attn==2.7.4.post1 --no-build-isolation
+```
+
+Verify:
+
+```bash
+python -c "import torch, transformers, flash_attn, pynvml; print(torch.__version__, transformers.__version__, flash_attn.__version__)"
+# expected: 2.6.0 4.51.3 2.7.4.post1
+```
+
+### 3. Download the model weights
+
+```bash
+huggingface-cli download --resume-download Anoy123423123/MSA-4B --local-dir ckpt/MSA-4B
+```
+
+This is the path `scripts/run_benchmarks.sh` expects (`model_path=ckpt/MSA-4B`). ~9 GB.
+
+### 4. Benchmark data
+
+Nothing to do — benchmark files are fetched on first use from
+[`Anoy123423123/MSA-RAG-BENCHMARKS`](https://huggingface.co/datasets/Anoy123423123/MSA-RAG-BENCHMARKS)
+into `./data/<benchmark>/`. Each benchmark is two pickle files: `mdata_*.pkl` (memory corpus) and
+`qdata_*.pkl` (queries). ~0.6 GB for everything; `ms_100M` alone is 440 MB.
+
+### 5. Run inference
+
+```bash
+bash scripts/run_benchmarks.sh eval_benchmark
+```
+
+Results land in `src/evaluation/outputs/eval_benchmark/` — one `.log` and one `.json` per benchmark —
+and retrieval precision metrics are printed as each benchmark finishes.
+
+Check the top of `scripts/run_benchmarks.sh` first:
+
+```bash
+export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7   # defaults to 8 GPUs
+model_path=ckpt/MSA-4B
+```
+
+The `benchmarks` array holds `name:batch_size` entries. **Batch size must be scaled together with the
+GPU count**: the memory corpus is partitioned into one bucket per visible GPU, so the defaults below
+assume 8 GPUs.
+
+| Category | Benchmark | Default batch size |
+|---|---|---|
+| Multi-hop QA | `2wikimultihopqa`, `hotpotqa`, `musique` | 16 |
+| Single-hop QA | `nature_questions`, `triviaqa_06M` | 16 |
+| Single-hop QA (large corpus) | `triviaqa_10M` | 2 |
+| HippoRAG | `hipporag_narrative`, `hipporag_popqa` | 16 |
+| Passage retrieval / multilingual | `msmarco_v1`, `dureader` | 16 |
+| Length scaling (100M tokens) | `ms_100M` | 16 |
+
+`ms_100M` and `triviaqa_06M` are commented out by default — uncomment them to reproduce the
+length-scaling results. `ms_100M` is the 100M-token setting and is by far the heaviest run.
+
+### 6. Score the answers
+
+Answer quality is graded by an LLM judge served through OpenRouter:
+
+```bash
+export OPENROUTER_API_KEY=<your key>          # or edit scripts/calculate_llm_score.sh
+bash scripts/calculate_llm_score.sh eval_benchmark
+```
+
+The argument must match the log directory name from step 5.
+
+---
+
+## Project Structure
+
+```
+├── scripts/
+│   ├── run_benchmarks.sh          # Run inference on benchmarks
+│   ├── calculate_llm_score.sh     # LLM-based answer evaluation
+│   └── resave_model.sh            # Convert base model to MSA format
+└── src/
+    ├── msa/                       # Core MSA implementation
+    │   ├── configuration_msa.py   # MSA model configuration
+    │   ├── memory_sparse_attention.py  # MemorySparseAttention layer
+    │   ├── model.py               # MSAForCausalLM / MSAModel
+    │   └── generate.py            # Generation logic
+    ├── config/
+    │   └── memory_config.py       # GenerateConfig, ModelConfig, MemoryConfig
+    ├── evaluation/
+    │   └── llm_judge.py           # LLM-based evaluation metrics
+    ├── app/
+    │   └── benchmark.py           # Benchmark runner
+    ├── utils/                     # GPU workers, caching, templates, etc.
+    ├── msa_service.py             # Multi-GPU inference engine (MSAEngine)
+    ���── prefill.py                 # Stage 1 prefill worker
+    ├── benchmarks.py              # Benchmark registry & specs
+    └── types.py                   # Core type definitions
+```
+
+---
+
+## Troubleshooting
+
+| Symptom | Cause / fix |
+|---|---|
+| `KeyError` / `AttributeError` while loading the model | Wrong `transformers` version. It must be exactly `4.51.3`. |
+| `import pynvml` fails or NVML calls are missing | The `pynvml` module here comes from `nvidia-ml-py`. If the separate, deprecated `pynvml` distribution is also installed it shadows the module: `pip uninstall -y pynvml && pip install nvidia-ml-py==12.575.51`. |
+| CUDA OOM | Lower the per-benchmark `batch_size` in `scripts/run_benchmarks.sh`, or reduce the GPU count together with the batch size. |
+| `flash_attn` import error | The wheel must match your torch / Python / CUDA combination. Build from source with `--no-build-isolation`. |
+| Benchmark data download stalls | Set `export HF_ENDPOINT=https://hf-mirror.com` and retry; partial files resume. |
+
+---
+
+## Citation
+
+Citation details are withheld during anonymous peer review.
+
+## License
+
+MIT.

requirements.txt

@@ -0,0 +1,29 @@
+# Tested with Python 3.12.9 on CUDA 12.x.
+# NOTE: flash-attn is intentionally NOT listed here, it needs --no-build-isolation.
+#       Install it separately, see step 3 of QUICK_START.md.
+
+# --- Core model stack (versions are strict: MSA targets the transformers 4.51 API) ---
+torch==2.6.0
+torchvision==0.21.0
+transformers==4.51.3
+accelerate==1.0.1
+liger_kernel==0.5.10
+
+# --- Hub / data IO ---
+# huggingface_hub must stay <1.0 to remain compatible with transformers 4.51.3.
+huggingface_hub==0.31.4
+datasets==3.1.0
+lmdb==1.6.2
+tqdm==4.67.1
+
+# --- Numerics ---
+numpy==1.26.4
+pillow==11.2.1
+packaging==25.0
+
+# --- GPU monitoring: provides the `pynvml` module.
+#     Do NOT replace with the deprecated `pynvml` package, its import path differs.
+nvidia-ml-py==12.575.51
+
+# --- LLM-based answer scoring (scripts/calculate_llm_score.sh) ---
+openai==1.79.0

scripts/calculate_llm_score.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+export OPENROUTER_API_KEY="your_openrouter_key"
+exp_name=$1
+python src/evaluation/llm_judge.py src/evaluation/outputs/${exp_name}

scripts/resave_model.sh

@@ -0,0 +1,38 @@
+# 模型超参数
+export POOLING_KERNEL_SIZE=64
+export TOP_K_DOCS=16
+
+# trick
+export AUX_LOSS="false"
+
+export REWRITE_POSITION="true"
+export ROUTER_LAYER_IDX="18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35"
+# export ROUTER_LAYER_IDX="20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39"
+
+# special tokens
+export ADD_MEM_TOKENS="false"
+export NUM_PREFIX_TOKENS=64
+export NUM_TAIL_TOKENS=64
+export FIRST_ADD_MEM_TOKENS="false"
+export WARM_UP_MEM_TOKENS="false"
+
+
+# loss weight
+export LMLOSS_WEIGHT=1.0
+export REC_LOSS_WEIGHT=0.0
+export AUX_LOSS_WEIGHT=0.1
+export ANS_LOSS_WEIGHT=1.0
+export AUX_LOSS_METHOD="INFONCE_DECOUPLE"  # INFONCE, INFONCE_FOCAL, BCE, INFONCE_DECOUPLE, INFONCE_DECOUPLE_FOCAL
+export INFONCE_LOSS_TEMP=0.1   # 如果使用infonce的话
+
+# 结构搜索
+export DECOUPLE_ROUTER="true"
+export HEAD_REDUCE_METHOD="mean"
+export QUERY_REDUCE_METHOD="max"
+export CHUNK_REDUCE_METHOD="max"
+export DECOUPLE_POOLING_MODE="mean"
+
+ORIGIN_MODEL_PATH=$1
+SAVE_MODEL_PATH=$2
+
+python src/utils/resave_model.py $ORIGIN_MODEL_PATH $SAVE_MODEL_PATH

scripts/run_benchmarks.sh

@@ -0,0 +1,119 @@
+#!/bin/bash
+export MASTER_PORT=29509
+export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+# export CUDA_VISIBLE_DEVICES=0,1
+
+# model
+model_path=ckpt/MSA-4B
+
+# data
+# bench_name:batch_size
+# Adjust the batch size appropriately based on the number of GPUs.
+benchmarks=(
+    # "ms_100M:16"
+    "hipporag_narrative:16"
+    "nature_questions:16"
+    "2wikimultihopqa:16"
+    "hotpotqa:16"
+    "musique:16"
+    "hipporag_popqa:16"
+    # "triviaqa_06M:16"
+    "triviaqa_10M:2"
+    "dureader:16"
+    "msmarco_v1:16"
+)
+
+top_p=0.9
+temperature=0.0
+max_length=2048
+template=QWEN3_INSTRUCT_TEMPLATE
+
+# Create log directory (use first argument as custom name, otherwise use timestamp)
+log_name="${1:-eval_logs_$(date +%Y%m%d_%H%M%S)}"
+log_dir="./src/evaluation/outputs/${log_name}"
+if [ -d "$log_dir" ]; then
+    echo "Error: log directory already exists: $log_dir"
+    exit 1
+fi
+mkdir -p "$log_dir"
+
+# Statistics
+total_benchmarks=${#benchmarks[@]}
+current=0
+success_count=0
+fail_count=0
+failed_benchmarks=()
+
+echo "=========================================="
+echo "Start running all benchmark evaluations"
+echo "Total: ${total_benchmarks} benchmarks"
+echo "Log directory: $log_dir"
+echo "=========================================="
+echo ""
+
+# Run each benchmark
+for entry in "${benchmarks[@]}"; do
+    benchmark="${entry%%:*}"
+    batch_size="${entry##*:}"
+    current=$((current + 1))
+    echo "[$current/$total_benchmarks] Running: $benchmark (batch_size=$batch_size)"
+    echo "Start time: $(date '+%Y-%m-%d %H:%M:%S')"
+
+    # Create separate log file for each benchmark
+    log_file="$log_dir/${benchmark}.log"
+    json_file="$log_dir/${benchmark}.json"
+
+    # Run evaluation and record logs
+    python -u src/app/benchmark.py \
+        --benchmark "$benchmark" \
+        --model_path "$model_path" \
+        --top_p "$top_p" \
+        --temperature "$temperature" \
+        --max_length "$max_length" \
+        --template "$template" \
+        --output_file "$json_file" \
+        --max_batch_size "$batch_size"  \
+        --max_chunk_per_block 16384 \
+        --block_size 2048 \
+        2>&1 | tee $log_file
+
+    # Check exit status
+    exit_code=${PIPESTATUS[0]}
+
+    if [ $exit_code -eq 0 ]; then
+        echo "[$current/$total_benchmarks] $benchmark finished (success)"
+        # Print benchmark name and metrics
+        echo "========== $benchmark Results =========="
+        python -c "import json; d=json.load(open('$json_file')); [print(f'  {k}: {v}') for k,v in d.get(list(d.keys())[0],{}).get('precision',{}).get('metrics',{}).items()]" 2>/dev/null || echo "  (failed to parse metrics)"
+        echo "========================================"
+        success_count=$((success_count + 1))
+    else
+        echo "[$current/$total_benchmarks] $benchmark failed (exit code: $exit_code)"
+        fail_count=$((fail_count + 1))
+        failed_benchmarks+=("$benchmark")
+    fi
+
+    echo "End time: $(date '+%Y-%m-%d %H:%M:%S')"
+    echo "----------------------------------------"
+    echo ""
+done
+
+# Summary
+echo "=========================================="
+echo "All benchmark evaluations completed"
+echo "=========================================="
+echo "Total: $total_benchmarks"
+echo "Success: $success_count"
+echo "Failed: $fail_count"
+echo ""
+
+if [ $fail_count -gt 0 ]; then
+    echo "Failed benchmarks:"
+    for failed in "${failed_benchmarks[@]}"; do
+        echo "  - $failed"
+    done
+    echo ""
+fi
+
+echo "All logs saved in: $log_dir"
+echo "=========================================="

src/app/benchmark.py

@@ -0,0 +1,335 @@
+import argparse
+from collections import OrderedDict
+import pickle
+import re
+from typing import Dict, List
+from tqdm import tqdm
+import json
+import torch
+import numpy as np
+import pathlib
+import sys
+import os
+import random
+import numpy as np
+import multiprocessing as mp
+
+project_path = pathlib.Path(__file__).parent.parent.parent
+sys.path.append(str(project_path))
+
+from src.benchmarks import BenchMarks
+
+
+def set_seed(seed):
+    """
+    固定所有随机种子以确保实验的可复现性。
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)  # 适用于多GPU环境
+    os.environ['PYTHONHASHSEED'] = str(seed)
+
+# 在程序的开始部分调用此函数
+seed_value = 42
+set_seed(seed_value)
+
+def parse_benchmark_file(args):
+    benchmark = BenchMarks(bench_name=args.benchmark)
+    path, mem_path = benchmark.get_bench_files()
+    args.query_file = path
+    args.memory_file = mem_path
+    print(" ==========> load query file from:", path)
+    if path.endswith(".json"):
+        with open(path, "r") as f:
+            raw_data = json.load(f)
+        ## 重组数据集
+        data = [{
+            'question': d['question'],
+            'labels': d['labels'],
+            'answer': d['answer']
+        } for d in raw_data]
+    # 判断是否是一个路径
+    elif path.endswith("pkl"):
+        with open(path, "rb") as f:
+            query_metas = pickle.load(f)
+
+        data = [{
+            'question': q_meta['query'],
+            'labels': q_meta['reference_list'],
+            'answer': q_meta['answer']
+        } for q_meta  in query_metas]
+    else:
+        raise ValueError(f"Unsupported file format: {path}")
+    print(f"num sample: {len(data)}")
+    return data
+
+def sort_requests(data_items: List[Dict]) -> List[dict]:
+    """
+    创建动态批次：按照input_id长度排序，然后根据max_input_length分批
+    """
+    print("Tokenizing and sorting data for dynamic batching...")
+
+    # 计算每个item的input_id长度
+    items_with_length = []
+    for item in tqdm(data_items):
+        prompt = item["question"]
+        length = len(prompt)
+
+        items_with_length.append((item, length))
+
+    # 按长度排序
+    items_with_length.sort(key=lambda x: x[1])
+    return [item[0] for item in items_with_length]
+
+
+def base_it(predict, label, at, score_func):
+    assert len(predict) == len(label)
+    scores = []
+    for pred, lbs in zip(predict, label):
+        pred = pred.tolist() if not isinstance(pred, list) else pred
+        best_score = 0.
+        if not isinstance(lbs, list):
+            lbs = [lbs]
+        for lb in lbs:
+            if isinstance(lb, list):
+                lb = lb[0]
+            rank = pred[:at].index(lb) + 1 if lb in pred[:at] else 0
+            cur_score = score_func(rank)
+            best_score = max(best_score, cur_score)
+        scores.append(best_score)
+    return scores
+
+
+def eval_recall(predict, label, at=10):
+    scores = base_it(predict, label, at, lambda rank: int(rank != 0))
+    return {f'R@{at}': sum(scores) / len(scores)}
+
+
+def eval_mrr(predict, label, at=10):
+    scores = base_it(predict, label, at, lambda rank: 1 / rank if rank != 0 else 0)
+    return {f'MRR@{at}': sum(scores) / len(scores)}
+
+
+def eval_all(predict, label):
+    log_dict = {}
+    log_dict.update(eval_recall(predict, label, at=1))
+    log_dict.update(eval_recall(predict, label, at=5))
+    log_dict.update(eval_recall(predict, label, at=10))
+    log_dict.update(eval_mrr(predict, label, at=1))
+    return log_dict
+
+def calculate_ir_metrics(true_labels: List[int], pred_labels: List[int]):
+    """计算信息检索中的 Precision, Recall, F1, 和 IoU。"""
+    true_set = set(true_labels)
+    pred_set = set(pred_labels)
+    if not true_set: return {'precision': 0.0, 'recall': 0.0, 'f1': 0.0, 'iou': 0.0}
+    tp = len(true_set & pred_set)
+    precision = tp / len(pred_set) if pred_set else 0.0
+    recall = tp / len(true_set)
+    f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0.0
+    iou = len(true_set & pred_set) / len(true_set | pred_set) if len(true_set | pred_set) > 0 else 0.0
+
+    return {'precision': precision, 'recall': recall, 'f1': f1, 'iou': iou}
+
+def process_results(requests: List[Dict], index_to_doc, doc_to_index):
+    num_error = 0
+    record_list = []
+    all_metrics = {"precision": [], "recall": [], "f1": [], "iou": []}
+
+    for request in requests:
+        question = request["question"]
+        answer = request["answer"]
+        generated_text = request["response"].replace('<|endoftext|>', '')
+        generated_text = "\nPlease answer the question based" + generated_text.split("\nPlease answer the question based")[1]
+
+        labels = [doc_to_index[txt] for txt in request["labels"]]
+        predictions = list(set(map(int, re.findall(r'\[(\d+)\]', generated_text))))
+
+        try:
+            pred_answer = generated_text.split('The answer to the question is:')[-1].split("<|im_end|>")[0].strip()
+        except Exception as e:
+            pred_answer = ""
+            raise ValueError(f"输出格式异常: {e}, generated_text: {generated_text}")
+
+        try:
+            record_list.append({
+                "labels_id": labels,
+                "pred_id": predictions,
+                "question": question,
+                "true_answer": answer,
+                "pred_answer": pred_answer,
+                "generated_text": generated_text,
+                "predict_context": [{i: index_to_doc[pid]} for i, pid in enumerate(predictions)],
+                "gt_context": [{i: index_to_doc[pid]} for i, pid in enumerate(labels)],
+            })
+        except Exception as e:
+            num_error += 1
+
+        metrics = calculate_ir_metrics(labels, predictions)
+        for k, v in metrics.items():
+            all_metrics[k].append(v)
+
+    metrics_dict = {k: round(float(np.mean(v)), 4) for k, v in all_metrics.items()}
+    print(" ==================== Retrieve metrics ======================= ")
+    print("AR Metrics: ", metrics_dict)
+    print(" ==================== Retrieve metrics ======================= ")
+    return {"metrics": metrics_dict, "record_list": record_list}
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--benchmark', type=str)
+    parser.add_argument('--max_batch_size', type=int)
+    parser.add_argument('--model_path', type=str)
+    parser.add_argument('--top_p', type=float, default=0.9)
+    parser.add_argument('--temperature', type=float, default=0.0)
+    parser.add_argument('--block_size', type=int, default=2048) # tokens for one memory inference
+    parser.add_argument('--max_chunk_per_block', type=int, default=16*1024) # chunks per block slice
+    parser.add_argument('--max_length', type=int, default=64) # max output length
+    parser.add_argument('--max_seq_len', type=int, default=0) # max input+output length
+    parser.add_argument('--max_query_seq_len', type=int, default=0) # max input seq len
+    parser.add_argument('--template', type=str, default="QWEN3_TEMPLATE")
+    parser.add_argument('--output_file', type=str, default="")
+    parser.add_argument('--case_name', type=str, default="anonymous")
+
+    args = parser.parse_args()
+
+    return args
+
+
+def should_regenerate(request:dict, response: str):
+    """return a text if the response should be regenerated, or return None"""
+
+    if response[-len("<|object_ref_end|>"):] == "<|object_ref_end|>" and "The answer to the question is:" not in response:
+        if not request.get('regenerated', False):
+            request["regenerated"] = True
+        response = "\nPlease answer the question based"+response.split("\nPlease answer the question based")[1]
+        response = response.replace('<|endoftext|>', '')
+        return "<regenerate>"+ response
+
+    return None
+
+def read_config_to_args(config_path):
+    with open(os.path.join(config_path, "config.json"), 'r') as f:
+        msa_config = json.load(f).get("msa_config")
+    args.doc_top_k = msa_config.get("doc_top_k", 16)
+    args.pooling_kernel_size = msa_config.get("pooling_kernel_size", 64)
+    args.router_layer_idx = msa_config.get("router_layer_idx", "all")
+    return args, msa_config
+
+
+def msa_benchmark(args, data):
+
+    from src.msa_service import GenerateConfig, ModelConfig, MemoryConfig, MSAEngine
+    args, msa_config = read_config_to_args(args.model_path)
+
+    model_config = ModelConfig(model_path=args.model_path,
+                                doc_top_k=args.doc_top_k,
+                                pooling_kernel_size=args.pooling_kernel_size,
+                                router_layer_idx=args.router_layer_idx,
+                                )
+    generate_config = GenerateConfig(devices=list(range(torch.cuda.device_count())),
+                                     template=args.template,
+                                     max_generate_tokens=args.max_length,
+                                     max_seq_len=args.max_seq_len,
+                                     max_query_seq_len=args.max_query_seq_len,
+                                     max_batch_size=args.max_batch_size,
+                                     top_p=args.top_p,
+                                     temperature=args.temperature,
+                                     qa_mode=True)
+    memory_config = MemoryConfig(block_size=args.block_size,
+                                 pooling_kernel_size=args.pooling_kernel_size,
+                                 slice_chunk_size=args.max_chunk_per_block,
+                                 memory_file_path=args.memory_file,
+                                 )
+    
+    final_result = {}
+    if args.output_file:
+        try:
+            with open(args.output_file, 'r') as f:
+                exist_result = json.load(f)
+                final_result = exist_result[args.case_name]
+        except:
+            pass
+
+    with MSAEngine(generate_config, model_config, memory_config) as engine:
+        print("start precision test")
+        idx_to_doc = engine.get_idx_to_doc()
+        doc_to_idx = {v: k for k, v in idx_to_doc.items()}
+
+        bsz = args.max_batch_size * generate_config.world
+        sorted_requests = sort_requests(data)
+        requests = OrderedDict({idx: item for idx, item in enumerate(sorted_requests)})
+
+        for req_idx in requests:
+            requests[req_idx]['idx'] = req_idx
+
+        total = len(requests)
+        results = [] # processed requests
+        pbar = tqdm(total=total, desc=f"Precision Test")
+
+        to_send = {} 
+        while len(results) < total:
+            num = min(bsz-len(to_send), len(requests))
+            for _ in range(num):
+                idx, request = requests.popitem(last=False)
+                to_send[idx] = request
+            
+            prompts, indices = [], [] 
+            for idx, request in to_send.items():
+                prompts.append(request.get("new_question", request["question"]))
+                indices.append(idx)
+
+            texts, recall_topks, _ = engine.generate(prompts, require_recall_topk=True)
+
+            for idx, response in enumerate(texts):
+                req_idx = indices[idx]
+                request = to_send[req_idx]
+                response = "\nPlease answer the question based"+response.split("\nPlease answer the question based")[1]
+                response = response.replace('<|endoftext|>', '')
+                new_prompt = should_regenerate(request, response)
+
+                # 判断是否需要重新生成
+                if new_prompt is not None:
+                    request["new_question"] = new_prompt
+                else:
+                    recall_topk = {layer: v[idx] for layer, v in recall_topks.items()}
+                    request = to_send.pop(req_idx)
+                    request['recall_topk'] = recall_topk
+                    request["response"] = response
+                    results.append(request)
+            pbar.update(bsz - len(to_send))
+        pbar.close()
+
+        assert len(results) == total, \
+            f"Results count mismatch: got {len(results)}, expected {total} (from query_file)"
+
+        final_result['precision'] = process_results(results, idx_to_doc, doc_to_idx)
+
+    if args.output_file:
+        exist_result = {}
+        try:
+            with open(args.output_file, 'r') as f:
+                exist_result = json.load(f)
+        except:
+            pass
+        exist_result[args.case_name] = final_result
+        with open(args.output_file, 'w') as f:
+            json.dump(exist_result, f, indent=4, ensure_ascii=False)
+    else:
+        s = json.dumps(final_result, indent=4, ensure_ascii=False)
+        print(s)
+
+if __name__ == "__main__":
+    mp.set_start_method('spawn') 
+    args = parse_args()
+    assert args.template in ["QWEN3_TEMPLATE", "QWEN3_INSTRUCT_TEMPLATE"]
+    if args.output_file:
+        assert args.case_name != "", "when output result to a file, please give this test case a name"
+
+    print(json.dumps(vars(args), indent=4, sort_keys=True))
+
+    data = parse_benchmark_file(args)
+    msa_benchmark(args, data)

src/benchmarks.py

@@ -0,0 +1,120 @@
+import os
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import ClassVar
+
+from huggingface_hub import hf_hub_download
+
+
+class Category(Enum):
+    """Benchmark categories with associated root directories and path patterns."""
+    RAG = auto()
+    RAG_0108 = auto()
+    LENGTH_SCALE = auto()
+
+
+# ============================================================================
+# HuggingFace config & local data root
+# ============================================================================
+
+HF_REPO_ID = "Anoy123423123/MSA-RAG-BENCHMARKS"
+_DATA_ROOT = os.path.join(os.getcwd(), "data")
+
+
+@dataclass(frozen=True)
+class BenchmarkSpec:
+    """Immutable specification for a single benchmark's file layout."""
+    bench_name: str       # benchmark name, also the HF subdirectory
+    query_file: str
+    memory_file: str
+
+    def _resolve(self, filename: str) -> str:
+        """Return local path if cached, otherwise download from HF into data/."""
+        local_path = os.path.join(_DATA_ROOT, self.bench_name, filename)
+        if os.path.exists(local_path):
+            return local_path
+        os.makedirs(os.path.dirname(local_path), exist_ok=True)
+        return hf_hub_download(
+            repo_id=HF_REPO_ID,
+            filename=f"{self.bench_name}/{filename}",
+            repo_type="dataset",
+            local_dir=_DATA_ROOT,
+        )
+
+    @property
+    def query_path(self) -> str:
+        return self._resolve(self.query_file)
+
+    @property
+    def memory_path(self) -> str:
+        return self._resolve(self.memory_file)
+
+    def get_bench_files(self) -> tuple[str, str]:
+        return self.query_path, self.memory_path
+
+
+# ============================================================================
+# Registry: benchmark name -> spec
+# ============================================================================
+
+def _rag(name: str) -> BenchmarkSpec:
+    return BenchmarkSpec(name, f"qdata_{name}.pkl", f"mdata_{name}.pkl")
+
+
+def _rag_0108(name: str) -> BenchmarkSpec:
+    return BenchmarkSpec(name, f"qdata_{name}.pkl", f"mdata_{name}.pkl")
+
+
+_REGISTRY: dict[str, BenchmarkSpec] = {
+    # --- Length-scale benchmarks ---
+    "ms_100M": BenchmarkSpec("ms_100M", "qdata_msmarco_16K.pkl", "mdata_msmarco_100M.pkl"),
+    # --- Multi-hop QA ---
+    "2wikimultihopqa":    _rag("2wikimultihopqa"),
+    "hotpotqa":           _rag("hotpotqa"),
+    "musique":            _rag("musique"),
+    # --- HippoRAG ---
+    "hipporag_narrative": _rag_0108("hipporag_narrative"),
+    "hipporag_popqa":     _rag_0108("hipporag_popqa"),
+    # --- Single-hop QA ---
+    "nature_questions":   _rag("nature_questions"),
+    "triviaqa_06M":       _rag("triviaqa_06M"),
+    "triviaqa_10M":       _rag("triviaqa_10M"),
+    # --- Multilingual / Passage retrieval ---
+    "dureader":           _rag("dureader"),
+    "msmarco_v1":         _rag("msmarco_v1"),
+}
+
+ALL_BENCH_NAMES: list[str] = list(_REGISTRY)
+
+
+# ============================================================================
+# Public API
+# ============================================================================
+
+class BenchMarks:
+    """Resolve benchmark name to query / memory file paths.
+
+    Usage:
+        bench = BenchMarks("hotpotqa")
+        query_file, memory_file = bench.get_bench_files()
+    """
+
+    AVAILABLE: ClassVar[list[str]] = ALL_BENCH_NAMES
+
+    def __init__(self, bench_name: str) -> None:
+        if bench_name not in _REGISTRY:
+            raise ValueError(
+                f"Unknown benchmark: {bench_name!r}. "
+                f"Available: {', '.join(ALL_BENCH_NAMES)}"
+            )
+        self._spec = _REGISTRY[bench_name]
+        self.name = bench_name
+        self.bench_name = self._spec.bench_name
+        self.query_file_name = self._spec.query_file
+        self.memory_file_name = self._spec.memory_file
+
+    def get_bench_files(self) -> tuple[str, str]:
+        return self._spec.get_bench_files()
+
+    def __repr__(self) -> str:
+        return f"BenchMarks({self.name!r})"

src/config/memory_config.py

@@ -0,0 +1,56 @@
+from typing import Union, Dict, List
+from dataclasses import dataclass
+
+from src.utils.template import QWEN3_TEMPLATE, QWEN3_INSTRUCT_TEMPLATE
+
+@dataclass
+class GenerateConfig:
+    devices: List[int] = None
+    template: Union[str, Dict] = None
+    max_generate_tokens: int = 256
+    max_seq_len: int = 0 # total sequence length in a batch
+    max_query_seq_len: int = 0 # max sequence for a single query
+    max_batch_size: int = 0 # 0 if batch size is not limited
+    top_p: float = 0.9
+    temperature: float = 0.0
+    qa_mode: bool = False
+
+    def __post_init__(self):
+        if isinstance(self.template, str):
+            assert self.template in ["QWEN3_TEMPLATE", "QWEN3_INSTRUCT_TEMPLATE"]
+            self.template = eval(self.template)
+        assert isinstance(self.template, dict)
+
+    @property
+    def world(self):
+        return len(self.devices) if self.devices else 0
+
+@dataclass
+class ModelConfig:
+    model_path: str = "Anoy123423123/MSA-4B"
+
+    doc_top_k: int = 16
+    pooling_kernel_size: int = 64
+    router_layer_idx: str = "all"
+
+    # template 
+    template_token_id = -2
+    template_id_num = 3
+
+    def get_model_envs(self):
+        envs = {}
+        # envs["TOP_K_DOCS"] = str(self.doc_top_k)
+        # envs["POOLING_KERNEL_SIZE"] = str(self.pooling_kernel_size)
+        # envs["ROUTER_LAYER_IDX"] = self.router_layer_idx
+        return envs
+
+
+@dataclass
+class MemoryConfig:
+    block_size: int = 16000  # 当对 memory 进行推理时使用的分块大小（tokens）
+    slice_chunk_size: int = 16 * 1024
+    pooling_kernel_size: int = 64
+    memory_file_path: str = ""
+
+    socket_ip: str = ""
+    socket_port: int = 0

src/evaluation/llm_judge.py

@@ -0,0 +1,218 @@
+from openai import OpenAI
+import os
+from tqdm import tqdm
+import sys
+import json
+import numpy as np
+from glob import glob
+from concurrent.futures import ThreadPoolExecutor
+
+def parse_match_result(text):
+    text_lower = text.strip().lower()
+    if 'match' in text_lower and 'mismatch' not in text_lower:
+        return 1
+    return 0
+
+def build_match_prompt(gold_answer, model_answer):
+    return f"""You are a strict but fair evaluator.
+
+Judge whether the Generated Answer correctly includes
+the core meaning of the Reference Answer.
+
+- The Generated Answer may add extra details.
+- Do NOT penalize additional information.
+- The core concept in the Reference Answer must be present
+  or clearly implied.
+
+  
+Reference Answer:
+{gold_answer}
+
+Generated Answer:
+{model_answer}
+
+Output:
+match or mismatch"""
+
+
+def build_score_prompt(gold_answer, model_answer, query):
+    return f""""Based on the accuracy, completeness, and relevance of the predicted answer to the real answer in the context of the **query**, assign an objective score from 0 to 5 (5 being the highest, 0 the lowest).
+
+    The scoring must strictly adhere to the following criteria. The final output can only be a single number.
+
+    Scoring Criteria:
+
+    5: The predicted answer is exactly the same as the real answer and correctly answers the query. Differences in wording do not affect factual accuracy.
+
+    4: The predicted answer contains all the core information of the real answer, with no errors, but includes a small amount of non-critical redundant content.
+
+    3: The predicted answer captures the core information but differs from the real answer in some aspects. The predicted answer is slightly incomplete or imprecise, but contains no errors.
+
+    2: The predicted answer is partially relevant to the real answer but omits a significant amount of information or deviates from the core topic of the query.
+
+    1: The predicted answer attempts to address the query (maintains basic relevance to the topic) but provides factually incorrect information. It does not contradict the core claim of the real answer, but shows incomplete or inaccurate understanding of the topic.
+
+    0. The predicted answer is completely unrelated to the query, consists of gibberish, or is a pure hallucination that shares no logical connection with the real answer.
+
+    Query:
+
+    {query}
+
+    True Answer:
+
+    {gold_answer}
+
+    Predicted Answer:
+
+    {model_answer}
+
+    Output only a single number (0, 1, 2, 3, 4, or 5): """
+
+
+
+def parse_score_result(text):
+    """解析 LLM 返回的评分结果，返回 0-5 的分数"""
+    text = text.strip()
+    # 尝试直接解析数字
+    for char in text:
+        if char.isdigit() and int(char) <= 5:
+            return int(char)
+    # 如果没有找到有效数字，返回 0
+    return 0
+
+def get_eval_response(prompt):
+    try:
+        completion = client.chat.completions.create(
+        extra_headers={
+        },
+        extra_body={},
+        model='google/gemini-2.5-flash',
+        messages=[
+            {
+            "role": "user",
+            "content": [
+                {
+                "type": "text",
+                "text":  prompt
+                },
+            ]
+            }
+        ],
+        temperature=0.0 
+        )
+        return completion.choices[0].message.content
+    except Exception as e:
+        print(e)
+        return ''
+
+if __name__ == "__main__":
+    dirname = sys.argv[1]
+    openrouter_api_key = os.environ.get("OPENROUTER_API_KEY", "")
+    if not openrouter_api_key:
+        raise ValueError("Please set the OPENROUTER_API_KEY environment variable.")
+    client = OpenAI(
+        base_url="https://openrouter.ai/api/v1",
+        api_key=openrouter_api_key,
+    )
+    JUDGE_MODEL = 'google/gemini-2.5-flash'
+    if dirname.endswith(".json") and os.path.isfile(dirname):
+        json_dir = os.path.split(dirname)[0]
+        json_paths = [dirname]
+    elif os.path.exists(dirname) and os.path.isdir(dirname):
+        json_dir = dirname
+        json_paths = glob(f'{json_dir}/*.json')
+    else:
+        json_dir = f"src/evaluation/outputs/{dirname}"
+        json_paths = glob(f'{json_dir}/*.json')
+
+    for json_path in tqdm(json_paths, total=len(json_paths)):
+        if 'score' in json_path:
+            continue
+        json_name = os.path.split(json_path)[-1].split('.')[0]
+        with open(json_path, 'r') as f:
+            datas = json.load(f)["anonymous"]["precision"]
+        print(datas["metrics"])
+
+        final_dict = dict()
+        final_dict['metrics'] = datas["metrics"]
+        final_dict['record_list'] = []
+        pred_answers = []
+        true_answers = []
+        questions    = []
+        pred_contexts = []
+        true_contexts = []
+        pred_ids = []
+        labels_ids = []
+
+        for dic in datas['record_list']:
+            question = dic['question']
+            questions.append(question)
+            true_answer = dic['true_answer']
+            key_word = "response" if "response" in dic else "pred_answer"
+            pred_answer = dic[key_word].replace("<|im_end|>",'').replace("<|endoftext|>",'')
+
+            pred_ids.append(dic['pred_id'])
+            labels_ids.append(dic['labels_id'])
+
+            pred_con_li = [v for di in dic['predict_context'] for k,v in di.items()]
+            true_con_li = [v for di in dic['gt_context'] for k,v in di.items()]
+            pred_contexts.append(pred_con_li)
+            true_contexts.append(true_con_li)
+
+            true_answers.append(true_answer)
+            qa = pred_answer.split('<answer>')[-1]
+            try:
+                q = qa.split("The answer to the question is: ")[-2].split("The user's question is: ")[-1].replace('\n<|object_ref_end|>','')
+            except:
+                q = ''
+            a = qa.split("The answer to the question is: ")[-1].replace('</answer>','').replace("The user's question is:",'')
+            a = a.split('</think>')[-1]
+            if "Answer:" in a:
+                a = a.split("Answer:")[1].strip()
+            pred_answers.append(a)
+        
+        match_prompts = []
+        score_prompts = []
+        for idx in tqdm(range(len(pred_answers))):
+            true_answer = true_answers[idx]
+            pred_answer = pred_answers[idx]
+            question = questions[idx]
+            match_prompt = build_match_prompt(true_answer, pred_answer)
+            score_prompt = build_score_prompt(true_answer, pred_answer, question)
+            match_prompts.append(match_prompt)
+            score_prompts.append(score_prompt)
+
+        
+        with ThreadPoolExecutor(max_workers=32) as executor:
+            score_results = list(tqdm(executor.map(get_eval_response, score_prompts), total=len(score_prompts), desc="Tokenizing texts"))
+
+        match_final_scores = []
+        score_final_scores = []
+        for idx in range(len(score_results)):
+            score_result = score_results[idx]
+            score_score = parse_score_result(score_result)
+            score_final_scores.append(score_score)
+
+        
+        match_final_scores_mean = np.mean(match_final_scores)
+        score_final_scores_mean = np.mean(score_final_scores)
+        print(" ======================================= ")
+        print(f"{json_name} LLM Score: {score_final_scores_mean:.4f}")
+        print(" ======================================= ")
+
+        old_file_path = json_path
+        new_file_path = json_path.replace(".json", f"_score{str(score_final_scores_mean).replace('.','point')[:9]}.json")
+        os.rename(old_file_path, new_file_path)
+
+        # final_dict
+        for idx in range(len(questions)):
+            question    =  questions[idx]
+            answer      =  true_answers[idx]
+            pred_answer =  pred_answers[idx]
+            score       =  score_final_scores[idx]
+            final_dict['record_list'].append({'question':question, 'true_answer':answer, 'pred_answer':pred_answer, 'pred_score':score})
+        final_dict["score_final_scores_mean"] = score_final_scores_mean
+        new_file_path = f'{json_dir}/{json_name}_llmscore.json'
+
+        with open(new_file_path, 'w',encoding='utf-8') as f:
+            json.dump(final_dict, f, indent=4,ensure_ascii=False)

src/msa/__init__.py

@@ -0,0 +1,3 @@
+from .generate import MSAGenerationMixin
+from .memory_sparse_attention import MemorySparseAttention
+from .configuration_msa import MSAConfig

src/msa/configuration_msa.py

@@ -0,0 +1,53 @@
+"""MSA (Memory Sparse Attention) Configuration"""
+
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+
+
+class DotDict(dict):
+    """支持点号访问的字典类"""
+    __getattr__ = dict.get
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+    
+    def __getstate__(self):
+        return dict(self)
+    
+    def __setstate__(self, state):
+        self.update(state)
+
+
+class MSAConfig(Qwen3Config):
+    """
+    MSA 模型的配置类，继承自 Qwen3Config。
+    
+    主要功能：确保 msa_config 在加载时自动转换为 DotDict，
+    支持使用点号访问属性（如 config.msa_config.pad_free）
+    """
+    
+    model_type = "msa"
+    
+    def __init__(self, msa_config=None, **kwargs):
+        super().__init__(**kwargs)
+        if msa_config is not None:
+            self.msa_config = DotDict(msa_config) if not isinstance(msa_config, DotDict) else msa_config
+    
+    def __setattr__(self, name, value):
+        """重写 __setattr__，确保设置 msa_config 时自动转换为 DotDict"""
+        if name == "msa_config" and isinstance(value, dict) and not isinstance(value, DotDict):
+            value = DotDict(value)
+        super().__setattr__(name, value)
+    
+    @classmethod
+    def from_dict(cls, config_dict, **kwargs):
+        """
+        从字典创建配置对象时，确保 msa_config 被转换为 DotDict。
+        这是关键方法，AutoConfig.from_pretrained() 最终会调用这个方法。
+        """
+        # 先调用父类的 from_dict
+        config = super().from_dict(config_dict, **kwargs)
+        
+        # 确保 msa_config 是 DotDict
+        if hasattr(config, 'msa_config') and isinstance(config.msa_config, dict) and not isinstance(config.msa_config, DotDict):
+            config.msa_config = DotDict(config.msa_config)
+        
+        return config

src/msa/generate.py

@@ -0,0 +1,354 @@
+import logging
+import os
+import re
+from typing import Optional, Union
+
+import torch
+from torch import nn
+from transformers.cache_utils import Cache
+from transformers.generation.configuration_utils import GenerationConfig
+from transformers.generation.logits_process import LogitsProcessorList
+from transformers.generation.stopping_criteria import StoppingCriteriaList
+from transformers.generation.streamers import BaseStreamer
+from transformers.generation.utils import (
+    GenerateDecoderOnlyOutput,
+    GenerateEncoderDecoderOutput,
+    GenerateNonBeamOutput,
+    GenerationMixin,
+)
+
+logger = logging.getLogger(__name__)
+
+class MSAGenerationMixin(GenerationMixin):
+    def _sample(
+        self,
+        input_ids: torch.LongTensor,
+        logits_processor: LogitsProcessorList,
+        stopping_criteria: StoppingCriteriaList,
+        generation_config: GenerationConfig,
+        synced_gpus: bool,
+        streamer: Optional["BaseStreamer"],
+        **model_kwargs,
+    ) -> Union[GenerateNonBeamOutput, torch.LongTensor]:
+        # init values
+        pad_token_id = generation_config._pad_token_tensor
+        output_attentions = generation_config.output_attentions
+        output_hidden_states = generation_config.output_hidden_states
+        output_scores = generation_config.output_scores
+        output_logits = generation_config.output_logits
+        return_dict_in_generate = generation_config.return_dict_in_generate
+        has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria)
+        do_sample = generation_config.do_sample
+
+        # init attention / hidden states / scores tuples
+        scores = () if (return_dict_in_generate and output_scores) else None
+        raw_logits = () if (return_dict_in_generate and output_logits) else None
+        decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+        cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+        decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+
+        # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+        if return_dict_in_generate and self.config.is_encoder_decoder:
+            encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
+            encoder_hidden_states = (
+                model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
+            )
+
+        # keep track of which sequences are already finished
+        batch_size, cur_len = input_ids.shape
+        this_peer_finished = False
+        unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
+        model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs)
+
+        model_forward = self.__call__
+        if isinstance(model_kwargs.get("past_key_values"), Cache):
+            is_compileable = model_kwargs["past_key_values"].is_compileable and self._supports_static_cache
+            if getattr(self, "hf_quantizer", None) is not None:
+                is_compileable &= self.hf_quantizer.is_compileable
+            is_compileable = is_compileable and not generation_config.disable_compile
+            if is_compileable and (
+                self.device.type == "cuda" or generation_config.compile_config._compile_all_devices
+            ):
+                os.environ["TOKENIZERS_PARALLELISM"] = "0"
+                model_forward = self.get_compiled_call(generation_config.compile_config)
+
+        if generation_config.prefill_chunk_size is not None:
+            model_kwargs = self._prefill_chunking(input_ids, generation_config, **model_kwargs)
+            is_prefill = False
+        else:
+            is_prefill = True
+
+        meta = model_kwargs["past_key_values"].meta
+        max_generate_tokens = model_kwargs["past_key_values"].meta["max_generate_tokens"]
+        tokenizer = meta['tokenizer']
+        response_string = meta["response_string"]
+        idx_to_doc = meta["idx_to_doc"]
+        pattern = meta["pattern"]
+        retrieval_end_flags = torch.zeros(batch_size, dtype=torch.bool, device=input_ids.device)
+        round_end_flags = torch.zeros(batch_size, dtype=torch.bool, device=input_ids.device)
+        inner_string = ["" for _ in range(batch_size)]
+        source_context_copied = False
+        all_input_str = [""] * batch_size
+        last_valid_inputs = input_ids[:, -1:].clone().to(input_ids.device)
+        is_first = 1
+        generate_stage = 1
+        cnt = 0
+        all_model_inputs = {}
+        has_generate_stage3 = False
+        first_stage2 = True
+        round_end = False
+
+        while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
+            if "position_ids" in model_kwargs:
+                position_ids = model_kwargs.pop("position_ids")
+            else:
+                position_ids = (position_ids[:, -1:] + model_kwargs["attention_mask"]).to(input_ids.device)
+
+            model_inputs = model_kwargs.copy()
+
+            if is_first == 1:
+                model_inputs.update({"input_ids": input_ids})
+            else:
+                model_inputs.update({"input_ids": last_valid_inputs})
+
+            model_inputs.update({"position_ids": position_ids})
+            model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
+            model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
+
+            input_str = tokenizer.batch_decode(model_inputs['input_ids'])
+            all_input_str = [s + input_str[i] for i, s in enumerate(all_input_str)]
+
+            if not is_prefill:
+                for layer_idx in range(self.config.num_hidden_layers):
+                    model_inputs["past_key_values"].record_kwargs(layer_idx, {"stage": "generate"})
+
+            if 'doc_ids' not in all_model_inputs:
+                all_model_inputs['doc_ids'] = model_inputs['doc_ids'].clone().to(input_ids.device)
+            else:
+                all_model_inputs['doc_ids'] = torch.cat((all_model_inputs['doc_ids'], model_inputs['doc_ids'][:, -model_inputs['attention_mask'].shape[1]:]), dim=1)
+
+            all_model_inputs['attention_mask'] = torch.cat((all_model_inputs['attention_mask'], model_inputs['attention_mask']), dim=1) if 'attention_mask' in all_model_inputs else model_inputs['attention_mask'].clone().to(input_ids.device)
+            all_model_inputs['input_ids'] = torch.cat((all_model_inputs['input_ids'], model_inputs['input_ids']), dim=1) if 'input_ids' in all_model_inputs else model_inputs['input_ids'].clone().to(input_ids.device)
+            all_model_inputs['position_ids'] = torch.cat((all_model_inputs['position_ids'], model_inputs['position_ids']), dim=1) if 'position_ids' in all_model_inputs else model_inputs['position_ids'].clone().to(input_ids.device)
+            all_model_inputs['past_key_values'] = model_inputs.get('past_key_values')
+            all_model_inputs['cache_position'] = model_inputs['cache_position'].clone().to(input_ids.device)
+
+            if round_end:
+                break
+
+            if is_prefill:
+                outputs = self(**all_model_inputs, return_dict=True)
+                is_prefill = False
+                first_stage2 = False
+            else:
+                outputs = model_forward(**model_inputs, return_dict=True)
+            is_first = 0
+
+            # update model kwargs for next generation step
+            model_kwargs = self._update_model_kwargs_for_generation(
+                outputs,
+                model_kwargs,
+                is_encoder_decoder=self.config.is_encoder_decoder,
+            )
+            if synced_gpus and this_peer_finished:
+                continue
+
+            next_token_logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device)
+
+            next_token_scores = logits_processor(last_valid_inputs, next_token_logits)
+
+            # Store scores, attentions and hidden_states when required
+            if return_dict_in_generate:
+                if output_scores:
+                    scores += (next_token_scores,)
+                if output_logits:
+                    raw_logits += (next_token_logits,)
+                if output_attentions:
+                    decoder_attentions += (
+                        (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
+                    )
+                    if self.config.is_encoder_decoder:
+                        cross_attentions += (outputs.cross_attentions,)
+                if output_hidden_states:
+                    decoder_hidden_states += (
+                        (outputs.decoder_hidden_states,)
+                        if self.config.is_encoder_decoder
+                        else (outputs.hidden_states,)
+                    )
+
+            # token selection
+            if do_sample:
+                probs = nn.functional.softmax(next_token_scores, dim=-1)
+                next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+            else:
+                next_tokens = torch.argmax(next_token_scores, dim=-1)
+            if has_eos_stopping_criteria:
+                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
+
+            cur_generate_context = tokenizer.batch_decode(next_tokens)
+
+            if generate_stage == 2 or generate_stage == 3:
+                source_context_list = []
+                imstart_str = "<|im_start|>"
+
+                if generate_stage == 2:
+                    for i, response_sample in enumerate(inner_string):
+                        if '<End-of-Retrieve>' in response_sample:
+                            question = all_input_str[i].split('historical document information\n\n')[-1]
+                            question = question.split('\nPlease return all documents related to the question')[0]
+                            source_context_list.append(imstart_str + 'The user\'s question is: %s\n<|object_ref_end|>' % (question))
+                        else:
+                            result = re.findall(pattern, response_sample)
+                            indices = sorted(list(set(map(int, result))))
+                            indices = [idx for idx in indices if idx in idx_to_doc]
+                            try:
+                                response_doc_str = ''.join(f"[{idx}]. {idx_to_doc[idx]}\n" for idx in indices)
+                                response_doc_str = response_doc_str + '<|object_ref_end|>'
+                            except KeyError as e:
+                                logger.warning("Document not found for index %s, available indices: %s", e, indices)
+                                response_doc_str = "" + '<|object_ref_end|>'
+                            source_context_list.append(response_doc_str)
+
+                if generate_stage == 3:
+                    for i, response_sample in enumerate(response_string):
+                        if '<End-of-Retrieve>' in response_sample:
+                            question = all_input_str[i].split('historical document information\n\n')[-1]
+                            question = question.split('\nPlease return all documents related to the question')[0]
+                            source_context_list.append(imstart_str + 'The user\'s question is: %s\n<|object_ref_end|>' % (question))
+                        else:
+                            source_context_list.append("")
+
+                source_batch = tokenizer(
+                    source_context_list,
+                    padding="longest",
+                    truncation=True,
+                    return_tensors="pt",
+                    add_special_tokens=True,
+                    padding_side="left",
+                )
+
+                sh = source_batch['input_ids'].shape
+
+                batch_source_input_ids = source_batch['input_ids'].clone().detach().long().to(input_ids.device)
+                batch_source_attn_mask = source_batch['attention_mask'].clone().detach().long().to(input_ids.device)
+                batch_source_doc_ids = torch.zeros_like(source_batch['input_ids'], dtype=torch.long, device=input_ids.device)
+                batch_source_position_ids = torch.arange(sh[1], dtype=torch.long, device=input_ids.device).unsqueeze(0).expand(sh[0], -1)
+                batch_source_position_ids = batch_source_position_ids + model_inputs['position_ids'] + torch.sum(batch_source_attn_mask, dim=1, keepdim=True) - sh[1] + 1
+
+                input_ids = batch_source_input_ids
+                model_kwargs['attention_mask'] = batch_source_attn_mask
+                model_kwargs['doc_ids'] = batch_source_doc_ids
+                model_kwargs['position_ids'] = batch_source_position_ids
+                source_context_copied = True
+
+                cur_len += sh[1]
+                this_peer_finished = False
+                del outputs
+                is_first = 1
+                inner_string = ["" for _ in range(batch_size)]
+
+                if generate_stage == 2:
+                    round_end = True
+
+                if generate_stage == 2:
+                    generate_stage = 1
+
+                if generate_stage == 3:
+                    generate_stage = 4
+
+                round_end_flags = torch.zeros(batch_size, dtype=torch.bool, device=input_ids.device)
+
+            else:
+                input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
+                if True:
+                    temp = stopping_criteria(input_ids[:, -1:], scores)
+                    unfinished_sequences = unfinished_sequences & ~temp
+                    # max_res_len = 3000
+                    # if max([len(n) for n in response_string]) > max_res_len:
+                    #     unfinished_sequences = torch.zeros_like(unfinished_sequences)
+
+                this_peer_finished = unfinished_sequences.max() == 0
+                if this_peer_finished and not source_context_copied:
+                    logger.warning("Generation finished before source context was copied")
+
+                if streamer is not None:
+                    streamer.put(next_tokens.cpu())
+                cur_len += 1
+                del outputs
+                input_ids = input_ids[:, -1:]
+                model_kwargs['doc_ids'] = torch.nn.functional.pad(model_kwargs['doc_ids'], (0, 1), value=-1)
+
+                model_kwargs["attention_mask"] = torch.ones(batch_size, 1, dtype=torch.long, device=input_ids.device)
+                if generate_stage == 1:
+                    for i, retrieval_end_flag in enumerate(retrieval_end_flags):
+                        if retrieval_end_flag or round_end_flags[i]:
+                            model_kwargs["attention_mask"][i, 0] = 0
+                        else:
+                            last_valid_inputs[i, -1] = input_ids[i, -1]
+                else:
+                    for i, retrieval_end_flag in enumerate(retrieval_end_flags):
+                        last_valid_inputs[i, -1] = input_ids[i, -1]
+
+                assert len(cur_generate_context) == len(response_string)
+                for i in range(len(cur_generate_context)):
+                    response_string[i] += cur_generate_context[i]
+                    inner_string[i] += cur_generate_context[i]
+                    round_end_flags[i] |= "<|object_ref_end|>" in inner_string[i]
+                    retrieval_end_flags[i] |= '<End-of-Retrieve>' in response_string[i]
+                    if source_context_copied and generate_stage == 1:
+                        mypattern = r"^\[\d*\]?$"
+                        is_id = bool(re.fullmatch(mypattern, inner_string[i]))
+                        is_EOR = inner_string[i] == '<End-of-Retrieve>'[:len(inner_string[i])]
+                        round_end_flags[i] |= (not is_id and not is_EOR)
+
+                    if not source_context_copied:
+                        max_ret_len = 1000
+                        retrieval_end_flags[i] |= len(response_string[i]) > max_ret_len
+
+                if sum(round_end_flags * unfinished_sequences * (~retrieval_end_flags)) == sum(unfinished_sequences * (~retrieval_end_flags)) and not has_generate_stage3:
+                    generate_stage = 2
+                    for i in range(len(cur_generate_context)):
+                        if source_context_copied and round_end_flags[i]:
+                            if "<|object_ref_end|>" not in inner_string[i] or not bool(re.fullmatch(mypattern, inner_string[i].split("<|object_ref_end|>")[0])):
+                                retrieval_end_flags[i] = True
+                if sum(retrieval_end_flags * unfinished_sequences) == sum(unfinished_sequences) and not has_generate_stage3:
+                    generate_stage = 3
+                    has_generate_stage3 = True
+            cnt += 1
+            if cnt > max_generate_tokens:
+                break
+
+        all_model_inputs['attention_mask'], indices = all_model_inputs['attention_mask'].sort(dim=1, descending=False, stable=True)
+        all_model_inputs['input_ids'] = all_model_inputs['input_ids'].gather(dim=1, index=indices)
+        all_model_inputs['position_ids'] = all_model_inputs['position_ids'].gather(dim=1, index=indices)
+        all_model_inputs['doc_ids'] = all_model_inputs['doc_ids'].gather(dim=1, index=indices)
+
+        input_ids = all_model_inputs['input_ids']
+
+        if streamer is not None:
+            streamer.end()
+
+        if return_dict_in_generate:
+            if self.config.is_encoder_decoder:
+                return GenerateEncoderDecoderOutput(
+                    sequences=input_ids,
+                    scores=scores,
+                    logits=raw_logits,
+                    encoder_attentions=encoder_attentions,
+                    encoder_hidden_states=encoder_hidden_states,
+                    decoder_attentions=decoder_attentions,
+                    cross_attentions=cross_attentions,
+                    decoder_hidden_states=decoder_hidden_states,
+                    past_key_values=model_kwargs.get("past_key_values"),
+                )
+            else:
+                return GenerateDecoderOnlyOutput(
+                    sequences=input_ids,
+                    scores=scores,
+                    logits=raw_logits,
+                    attentions=decoder_attentions,
+                    hidden_states=decoder_hidden_states,
+                    past_key_values=model_kwargs.get("past_key_values"),
+                )
+        else:
+            return input_ids

src/msa/memory_sparse_attention.py

@@ -0,0 +1,852 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3Attention,
+    Qwen3Config,
+    apply_rotary_pos_emb,
+    repeat_kv,
+)
+try:
+    from flash_attn import flash_attn_varlen_func
+except ImportError:
+    print("请安装flash-attn库: pip install flash-attn --no-build-isolation")
+    flash_attn_varlen_func = None
+
+
+class MemorySparseAttention(Qwen3Attention):
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__(config=config, layer_idx=layer_idx)
+        if flash_attn_varlen_func is None:
+            raise ImportError("flash_attn is required. Please install it via 'pip install flash-attn --no-build-isolation'")
+        
+        self.layer_idx = layer_idx
+        self.top_k_docs = config.msa_config.top_k_docs
+        self.pooling_kernel_size = config.msa_config.pooling_kernel_size
+        self.router_layer_idx = config.msa_config.router_layer_idx
+
+        if self.router_layer_idx == "all":
+            self.router_layer_idx = list(range(config.num_hidden_layers))
+        else:
+            self.router_layer_idx = [int(i) for i in self.router_layer_idx.split(",")]
+        self.is_router_layer = self.layer_idx in self.router_layer_idx
+
+        self.head_reduce_method = config.msa_config.head_reduce_method
+        self.query_reduce_method = config.msa_config.query_reduce_method
+        self.chunk_reduce_method = config.msa_config.chunk_reduce_method
+        self.decouple_pooling_mode = config.msa_config.decouple_pooling_mode
+        self.aux_loss_method = config.msa_config.aux_loss_method
+
+        self.decouple_router = config.msa_config.decouple_router
+        if self.is_router_layer and self.decouple_router:
+            self.router_k_proj = nn.Sequential(
+                nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False),
+                # nn.GELU(),
+                # nn.Linear(config.num_key_value_heads * self.head_dim, config.num_key_value_heads * self.head_dim, bias=False)
+            )
+            self.router_q_proj = nn.Sequential(
+                nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False),
+                # nn.GELU(),
+                # nn.Linear(config.num_attention_heads * self.head_dim, config.num_attention_heads * self.head_dim, bias=False)
+            )
+        self.num_kv_heads = config.num_key_value_heads
+
+        self.sliding_window = None
+        self.selected_docs_indices = None
+        self.max_doc_id = None
+        self.num_split_for_kv = 8
+        self.template_prefix_kcache = None
+        self.template_prefix_vcache = None
+        self.memory_client = None
+
+    def set_memory_client(self, memory_client):
+        self.memory_client = memory_client
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        doc_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        
+        if self.training:
+            return self._forward(
+                hidden_states,
+                doc_ids,
+                attention_mask,
+                position_embeddings,
+                past_key_value,
+                **kwargs,
+            )
+        elif past_key_value is not None:
+            return self.forward_with_kvcache_for_batch_parrallel(
+                hidden_states,
+                doc_ids,
+                attention_mask,
+                position_embeddings,
+                past_key_value,
+                **kwargs,
+            )
+        else:
+            raise Exception("error!")
+
+    @staticmethod
+    def map_tensor_to_group_ids(a: torch.Tensor) -> torch.Tensor:
+        if a.ndim != 1:
+            raise ValueError("输入 Tensor a 必须是一维的。")
+
+        diff_mask = torch.diff(a) != 0  # [L-1]
+        id_increments = diff_mask.int()  # [L-1]
+        group_indices_offset = torch.cumsum(id_increments, dim=0)  # [L-1]
+
+        b = torch.cat((
+            torch.tensor([0], device=a.device, dtype=a.dtype), 
+            group_indices_offset
+        )) + 1
+        
+        return b
+
+    def forward_with_kvcache_for_batch_parrallel(
+        self,
+        hidden_states: torch.Tensor,
+        doc_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        
+        bsz, q_len, _ = hidden_states.shape
+        device, dtype = hidden_states.device, hidden_states.dtype
+        hidden_shape = (bsz, q_len, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        
+        stage = past_key_value.cache_kwargs[self.layer_idx]["stage"]
+
+        if stage == "prefill_stage1":
+            max_doc_id = int(doc_ids.max().item())
+            doc_token_mask = (doc_ids > 0) & (attention_mask == 1)
+            doc_indices = torch.nonzero(doc_token_mask, as_tuple=False)
+            original_doc_ids = doc_ids[doc_token_mask]
+            original_doc_batch_indices = doc_indices[:, 0]
+            global_doc_ids = original_doc_batch_indices * (max_doc_id + 1) + original_doc_ids
+            
+            if self.is_router_layer:
+                _, counts = torch.unique_consecutive(global_doc_ids, return_counts=True)
+                total_doc_tokens = global_doc_ids.shape[0]
+
+                cu_seqlens = counts.cumsum(0)
+                offsets = torch.zeros(counts.shape[0] + 1, dtype=counts.dtype, device=device)
+                offsets[1:] = cu_seqlens
+                offsets = offsets[:-1]
+
+                expanded_offsets = torch.repeat_interleave(offsets, counts)
+                original_order_ranks = torch.arange(total_doc_tokens, device=device) - expanded_offsets
+                
+                chunk_indices = original_order_ranks // self.pooling_kernel_size
+                max_chunks_per_doc = (q_len // self.pooling_kernel_size) + 1
+
+                global_chunk_ids = global_doc_ids * max_chunks_per_doc + chunk_indices
+
+                unique_global_chunk_ids, chunk_token_counts = torch.unique_consecutive(global_chunk_ids, return_counts=True)
+                pooled_doc_ids = unique_global_chunk_ids // max_chunks_per_doc % (max_doc_id + 1)
+
+
+                pooled_k_chunks, pooled_v_chunks = self.sequence_pooling_kv(
+                    key_states,
+                    value_states,
+                    doc_indices,
+                    global_chunk_ids,
+                )
+
+                pooled_k_chunks = pooled_k_chunks.transpose(0, 1).unsqueeze(0)
+                pooled_v_chunks = pooled_v_chunks.transpose(0, 1).unsqueeze(0)
+
+                pooled_router_k = None
+                if self.decouple_router:
+                    r_k_raw = self.router_k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+                    r_k_docs = r_k_raw[doc_indices[:, 0], :, doc_indices[:, 1]]
+                    
+                    _, chunk_lengths = torch.unique_consecutive(global_chunk_ids, return_counts=True)
+                    chunk_counts_view = chunk_lengths.view(-1, 1, 1).to(dtype=torch.float32)
+                    b_k, h_k, d_k = r_k_docs.shape
+                    k_flat = r_k_docs.reshape(b_k, -1).to(dtype=torch.float32)
+                    k_cumsum = F.pad(torch.cumsum(k_flat, dim=0), (0, 0, 1, 0))
+                    chunk_cu_seqlens = F.pad(torch.cumsum(chunk_lengths, 0), (1, 0))
+                    k_sums_flat = k_cumsum[chunk_cu_seqlens[1:]] - k_cumsum[chunk_cu_seqlens[:-1]]
+                    pooled_router_k = (k_sums_flat.view(unique_global_chunk_ids.shape[0], h_k, d_k) / chunk_counts_view).to(dtype=r_k_docs.dtype)
+                    
+                    pooled_router_k = pooled_router_k.transpose(0, 1).unsqueeze(0)
+                
+                if self.aux_loss_method == "INFONCE":
+                    router_k = pooled_router_k if pooled_router_k is not None else pooled_k_chunks
+                    pooled_router_k = F.normalize(router_k, p=2, dim=-1)
+
+            if past_key_value is not None:
+                num_template_mask_prefix = (doc_ids == -2).sum()
+                template_prefix_kcache = key_states[:, :, :num_template_mask_prefix]
+                template_prefix_vcache = value_states[:, :, :num_template_mask_prefix]
+                kwargs = {
+                    "template_prefix_kcache": template_prefix_kcache,
+                    "template_prefix_vcache": template_prefix_vcache,
+                }
+                if self.is_router_layer:
+                    pooled_k_chunks, pooled_v_chunks = past_key_value.update(pooled_k_chunks, pooled_v_chunks, self.layer_idx)
+                    kwargs2 = {
+                        "doc_id_bias": doc_ids.shape[1],
+                        "pooled_doc_ids": pooled_doc_ids,
+                        "prefill_stage1_kvcache_size": pooled_k_chunks.shape[2],
+                    }
+                    if pooled_router_k is not None:
+                        past_key_value.update_router_kcache(pooled_router_k, self.layer_idx)
+                    kwargs.update(kwargs2)
+                past_key_value.record_kwargs(self.layer_idx, kwargs)
+
+            key_states = repeat_kv(key_states, self.num_key_value_groups)
+            value_states = repeat_kv(value_states, self.num_key_value_groups)
+            
+            attn_output = torch.zeros((bsz, q_len, self.config.num_attention_heads * self.head_dim), device=device, dtype=dtype)
+            indices_b = torch.nonzero(doc_token_mask, as_tuple=False)
+            
+            if indices_b.shape[0] > 0:
+                q_b, k_b, v_b = query_states[indices_b[:, 0], :, indices_b[:, 1]], key_states[indices_b[:, 0], :, indices_b[:, 1]], value_states[indices_b[:, 0], :, indices_b[:, 1]]
+                doc_ids_b = doc_ids[indices_b[:, 0], indices_b[:, 1]]
+                batch_indices_b = indices_b[:, 0]
+                global_doc_ids_b = batch_indices_b * (max_doc_id + 1) + doc_ids_b
+                _, counts_b = torch.unique_consecutive(global_doc_ids_b, return_counts=True)
+                cu_seqlens_b = F.pad(torch.cumsum(counts_b, dim=0, dtype=torch.int32), (1, 0))
+                output_b_flat = flash_attn_varlen_func(q_b, k_b, v_b, cu_seqlens_q=cu_seqlens_b, cu_seqlens_k=cu_seqlens_b, max_seqlen_q=int(counts_b.max()), max_seqlen_k=int(counts_b.max()), dropout_p=self.attention_dropout if self.training else 0.0, causal=True).view(-1, self.config.num_attention_heads * self.head_dim)
+                attn_output[indices_b[:, 0], indices_b[:, 1]] += output_b_flat
+            
+            template_mask = (doc_ids == -2) & (attention_mask == 1)
+            template_indices = torch.nonzero(template_mask, as_tuple=False)
+            if template_indices.shape[0] > 0:
+                q_template = query_states.transpose(1, 2)[template_mask]
+                k_template = key_states.transpose(1, 2)[template_mask]
+                v_template = value_states.transpose(1, 2)[template_mask]
+                template_counts_per_sample = torch.bincount(template_indices[:, 0], minlength=bsz)
+                cu_seqlens_template = F.pad(torch.cumsum(template_counts_per_sample, dim=0, dtype=torch.int32), (1, 0))
+                output_template_flat = flash_attn_varlen_func(q_template, k_template, v_template, cu_seqlens_q=cu_seqlens_template, cu_seqlens_k=cu_seqlens_template, max_seqlen_q=int(template_counts_per_sample.max()), max_seqlen_k=int(template_counts_per_sample.max()), dropout_p=0.0, causal=True).view(-1, self.config.num_attention_heads * self.head_dim)
+                attn_output[template_mask] = output_template_flat
+
+            return self.o_proj(attn_output), None
+
+        elif stage == "prefill_stage2":
+            cache_kwargs = past_key_value.cache_kwargs[self.layer_idx]
+            if self.memory_client is not None:
+                if self.template_prefix_kcache is None:
+                    self.template_prefix_kcache , self.template_prefix_vcache  = self.memory_client.get_template_prefix_kvcaches(self.layer_idx)
+                    if not self.template_prefix_kcache.is_cuda:
+                        self.template_prefix_kcache = self.template_prefix_kcache.to(device)
+                    if not self.template_prefix_vcache.is_cuda:
+                        self.template_prefix_vcache = self.template_prefix_vcache.to(device)
+                template_prefix_kcache = self.template_prefix_kcache
+                template_prefix_vcache = self.template_prefix_vcache
+            else:
+                template_prefix_kcache = cache_kwargs["template_prefix_kcache"].to(device)
+                template_prefix_vcache = cache_kwargs["template_prefix_vcache"].to(device)
+
+            final_k_to_scatter, final_v_to_scatter = None, None
+
+            if self.is_router_layer:
+                routing_q_for_scoring = self.router_q_proj(hidden_states).view(hidden_shape).transpose(1, 2) if self.decouple_router else query_states
+                if self.aux_loss_method == "INFONCE":
+                    routing_q_for_scoring = F.normalize(routing_q_for_scoring, p=2, dim=-1)
+
+                query_mask = ((doc_ids == 0) & (attention_mask == 1))
+                res = self.memory_client.doc_query(routing_q_for_scoring, query_mask, self.layer_idx)
+                final_k_to_scatter, final_v_to_scatter, final_scores, num_selected_chunks_per_sample, final_selected_doc_ids = res
+
+                if past_key_value.meta.get("require_recall_topk", False):
+                    recall_topk_list = []
+                    for i in range(bsz):
+                        recall_topk_list.append({
+                            "topk_doc_ids": final_selected_doc_ids[i].cpu().detach().tolist(),
+                            "score": final_scores[i].cpu().detach().tolist(),
+                        })
+                    cache_kwargs["recall_topk"] = recall_topk_list
+            else:
+                num_selected_chunks_per_sample = torch.zeros(bsz, dtype=torch.long, device=device)
+
+            num_q_per_sample = attention_mask.sum(dim=1)
+            template_len = template_prefix_kcache.shape[2]
+            kv_lengths = template_len + num_selected_chunks_per_sample + num_q_per_sample
+
+            cu_seqlens_q = F.pad(num_q_per_sample.cumsum(0, dtype=torch.int32), (1, 0))
+            cu_seqlens_kv = F.pad(kv_lengths.cumsum(0, dtype=torch.int32), (1, 0))
+            
+            total_q_tokens = cu_seqlens_q[-1].item()
+            total_kv_tokens = cu_seqlens_kv[-1].item()
+
+            q_final = torch.empty((total_q_tokens, self.config.num_attention_heads, self.head_dim), device=device, dtype=dtype)
+            k_final_unrepeated = torch.empty((self.config.num_key_value_heads, total_kv_tokens, self.head_dim), device=device, dtype=dtype)
+            v_final_unrepeated = torch.empty((self.config.num_key_value_heads, total_kv_tokens, self.head_dim), device=device, dtype=dtype)
+
+            valid_q_mask = (attention_mask == 1)
+            q_final = query_states.permute(0, 2, 1, 3)[valid_q_mask]
+
+            offset_start_sample = cu_seqlens_kv[:-1]
+            offset_start_template = offset_start_sample
+            offset_start_chunks = offset_start_sample + template_len
+            offset_start_question = offset_start_chunks + num_selected_chunks_per_sample
+
+            template_indices = torch.arange(template_len, device=device).unsqueeze(0) + offset_start_template.unsqueeze(1)
+            source_k_template = template_prefix_kcache.expand(bsz, -1, -1, -1).permute(1, 0, 2, 3).reshape(self.config.num_key_value_heads, -1, self.head_dim)
+            k_final_unrepeated[:, template_indices.flatten(), :] = source_k_template
+            source_v_template = template_prefix_vcache.expand(bsz, -1, -1, -1).permute(1, 0, 2, 3).reshape(self.config.num_key_value_heads, -1, self.head_dim)
+            v_final_unrepeated[:, template_indices.flatten(), :] = source_v_template
+            
+            if self.is_router_layer and final_k_to_scatter is not None and final_k_to_scatter.shape[1] > 0:
+                batch_indices_for_chunks = torch.arange(bsz, device=device).repeat_interleave(num_selected_chunks_per_sample)
+                is_start_of_sample = torch.cat([torch.tensor([True], device=device), batch_indices_for_chunks[1:] != batch_indices_for_chunks[:-1]])
+                cumsum_ranks = torch.ones_like(batch_indices_for_chunks).cumsum(0)
+                start_offsets = cumsum_ranks[is_start_of_sample].repeat_interleave(num_selected_chunks_per_sample)
+                chunk_rank_in_sample = cumsum_ranks - start_offsets
+
+                chunk_dest_indices = offset_start_chunks[batch_indices_for_chunks] + chunk_rank_in_sample
+                
+                k_final_unrepeated[:, chunk_dest_indices, :] = final_k_to_scatter
+                if final_v_to_scatter.device == torch.device("cpu"):
+                    final_v_to_scatter = final_v_to_scatter.to(device)
+                v_final_unrepeated[:, chunk_dest_indices, :] = final_v_to_scatter
+
+            batch_indices_for_q = torch.arange(bsz, device=device).repeat_interleave(num_q_per_sample)
+            q_rank_in_sample = (torch.cumsum(valid_q_mask.int(), dim=1) - 1)[valid_q_mask]
+            q_dest_indices = offset_start_question[batch_indices_for_q] + q_rank_in_sample
+            
+            k_final_unrepeated[:, q_dest_indices, :] = key_states.permute(1, 0, 2, 3).reshape(self.config.num_key_value_heads, -1, self.head_dim)[:, valid_q_mask.flatten(), :]
+            v_final_unrepeated[:, q_dest_indices, :] = value_states.permute(1, 0, 2, 3).reshape(self.config.num_key_value_heads, -1, self.head_dim)[:, valid_q_mask.flatten(), :]
+
+            k_final = k_final_unrepeated
+            v_final = v_final_unrepeated
+
+            output_flat = flash_attn_varlen_func(
+                q=q_final, k=k_final.transpose(0,1), v=v_final.transpose(0,1),
+                cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_kv,
+                max_seqlen_q=num_q_per_sample.max().item(), max_seqlen_k=kv_lengths.max().item(),
+                dropout_p=0.0, causal=True
+            ).view(-1, self.config.num_attention_heads * self.head_dim)
+
+            attn_output = torch.zeros((bsz, q_len, self.config.num_attention_heads * self.head_dim), device=device, dtype=dtype)
+            attn_output[valid_q_mask] = output_flat
+            
+            max_kv_len = kv_lengths.max().item()
+            compacked_key_cache = torch.zeros((bsz, self.config.num_key_value_heads, max_kv_len, self.head_dim), dtype=dtype, device=device)
+            compacked_value_cache = torch.zeros((bsz, self.config.num_key_value_heads, max_kv_len, self.head_dim), dtype=dtype, device=device)
+
+            left_pad_mask = torch.arange(max_kv_len, device=device).unsqueeze(0) >= (max_kv_len - kv_lengths.unsqueeze(1))
+            
+            compacked_key_cache.permute(0, 2, 1, 3)[left_pad_mask] = k_final_unrepeated.permute(1, 0, 2)
+            compacked_value_cache.permute(0, 2, 1, 3)[left_pad_mask] = v_final_unrepeated.permute(1, 0, 2)
+
+            cache_kwargs["compacked_key_cache"] = compacked_key_cache
+            cache_kwargs["compacked_value_cache"] = compacked_value_cache
+            cache_kwargs["kv_lengths"] = kv_lengths
+            cache_kwargs["attention_mask"] = left_pad_mask
+            past_key_value.record_kwargs(self.layer_idx, cache_kwargs)
+            
+            return self.o_proj(attn_output), None
+
+        else: 
+            cache_kwargs = past_key_value.cache_kwargs[self.layer_idx]
+            if "compacked_key_cache" not in cache_kwargs:
+                raise ValueError("批次化紧凑KV缓存未找到。Prefill stage 2 是否正确运行?")
+
+            compacked_key_cache = cache_kwargs["compacked_key_cache"]
+            compacked_value_cache = cache_kwargs["compacked_value_cache"]
+            kv_lengths = cache_kwargs["kv_lengths"]
+            layer_attention_mask = cache_kwargs["attention_mask"]
+            
+            max_kv_len = compacked_key_cache.shape[2]
+            full_k_unrepeated = torch.cat([compacked_key_cache, key_states], dim=2)
+            full_v_unrepeated = torch.cat([compacked_value_cache, value_states], dim=2)
+
+            if past_key_value.meta.get("qa_mode", False):
+                cur_layer_attention_mask = torch.LongTensor([[1] * q_len for _ in range(bsz)]).to(device)
+                cur_layer_attention_mask = (cur_layer_attention_mask * attention_mask).type(layer_attention_mask.dtype)
+                layer_attention_mask = torch.cat([layer_attention_mask, cur_layer_attention_mask], dim=1)
+                attn_mask_4d = layer_attention_mask[:, None, None, :].expand(-1, self.config.num_attention_heads, 1, -1)
+                cache_kwargs["attention_mask"] = layer_attention_mask
+            else:
+                new_kv_lengths = kv_lengths + 1
+                max_new_kv_len = max_kv_len + 1
+                attn_mask_2d = torch.arange(max_new_kv_len, device=device).unsqueeze(0) >= (max_new_kv_len - new_kv_lengths.unsqueeze(1))
+
+                attn_mask_4d = attn_mask_2d[:, None, None, :].expand(-1, self.config.num_attention_heads, 1, -1)
+                cache_kwargs["kv_lengths"] = new_kv_lengths
+
+            key_states_gqa = repeat_kv(full_k_unrepeated, self.num_key_value_groups)
+            value_states_gqa = repeat_kv(full_v_unrepeated, self.num_key_value_groups)
+            
+            attn_output = F.scaled_dot_product_attention(
+                query_states, 
+                key_states_gqa, 
+                value_states_gqa, 
+                attn_mask=attn_mask_4d, 
+                dropout_p=0.0, 
+                is_causal=False
+            ).transpose(1, 2).reshape(bsz, q_len, -1)
+            
+            cache_kwargs["compacked_key_cache"] = full_k_unrepeated
+            cache_kwargs["compacked_value_cache"] = full_v_unrepeated
+            past_key_value.record_kwargs(self.layer_idx, cache_kwargs)
+            
+            return self.o_proj(attn_output), None
+
+    def _calculate_routing_scores_adaptive(
+        self,
+        query_states: torch.Tensor,     # [B, H, Q_len, D]
+        pooled_k_bched: torch.Tensor,   # [B, C, H, D]
+        routing_query_mask: torch.Tensor, # [B, Q_len] - 1 for valid, 0 for pad
+        chunk_mask: torch.Tensor,       # [B, C] - 1 for valid, 0 for pad
+    ) -> torch.Tensor:
+        bsz, num_heads, q_len, head_dim = query_states.shape
+        _, max_chunks, _, _ = pooled_k_bched.shape
+        dtype, device = query_states.dtype, query_states.device
+        min_val = torch.finfo(dtype).min
+
+        k_states_T = pooled_k_bched.permute(0, 2, 3, 1)
+
+        current_scaling = 1.0 if self.decouple_router and "INFONCE" in self.aux_loss_method else self.scaling
+        scores = torch.matmul(query_states, k_states_T) * current_scaling
+        
+        q_mask_expanded = routing_query_mask.view(bsz, 1, q_len, 1)
+        k_mask_expanded = chunk_mask.view(bsz, 1, 1, max_chunks)
+        
+        final_mask = q_mask_expanded & k_mask_expanded
+        scores.masked_fill_(~final_mask, min_val)
+
+        if self.head_reduce_method == "max":
+            scores = scores.max(dim=1).values
+        elif self.head_reduce_method == "mean":
+            scores = scores.mean(dim=1)
+        else:
+            raise NotImplementedError(f"Unsupported head reduce method: {self.head_reduce_method}")
+
+        if self.query_reduce_method == "max":
+            scores_final = scores.max(dim=1).values
+
+        elif self.query_reduce_method == "mean":
+            valid_mask = final_mask.squeeze(1) # [B, Q_len, C]
+
+            scores_clean = torch.where(valid_mask, scores, torch.zeros_like(scores))
+            sum_scores = scores_clean.sum(dim=1) # [B, C]
+            counts = valid_mask.sum(dim=1).to(dtype).clamp(min=1.0)
+            mean_scores = sum_scores / counts
+
+            scores_final = torch.where(
+                chunk_mask, 
+                mean_scores, 
+                torch.tensor(min_val, device=device, dtype=dtype)
+            )
+
+        elif self.query_reduce_method == "last":
+            q_lens = routing_query_mask.sum(dim=1).long()
+            last_indices = (q_lens - 1).clamp(min=0)
+
+            gather_idx = last_indices.view(bsz, 1, 1).expand(-1, 1, max_chunks)
+            scores_final = scores.gather(1, gather_idx).squeeze(1)
+            scores_final.masked_fill_(~chunk_mask, min_val)
+            
+        else:
+            raise NotImplementedError(f"Unsupported query reduce method: {self.query_reduce_method}")
+
+        return scores_final
+
+    def sequence_pooling_kv(self, key_states, value_states, doc_indices, global_chunk_ids):
+        k_docs = key_states[doc_indices[:, 0], :, doc_indices[:, 1]]
+        v_docs = value_states[doc_indices[:, 0], :, doc_indices[:, 1]]
+        unique_global_chunk_ids, chunk_lengths = torch.unique_consecutive(global_chunk_ids, return_counts=True)
+        
+        num_unique_chunks = unique_global_chunk_ids.shape[0]
+        chunk_counts_view = chunk_lengths.view(-1, 1, 1).to(dtype=torch.float32)
+
+        def compute_pooled_states_via_cumsum(states, counts_view, lengths):
+            b, h, d = states.shape
+            states_flat = states.reshape(b, -1).to(dtype=torch.float32)
+            states_cumsum = F.pad(torch.cumsum(states_flat, dim=0), (0, 0, 1, 0))
+            chunk_cu_seqlens = F.pad(torch.cumsum(lengths, 0), (1, 0))
+            state_sums_flat = states_cumsum[chunk_cu_seqlens[1:]] - states_cumsum[chunk_cu_seqlens[:-1]]
+            state_sums = state_sums_flat.view(num_unique_chunks, h, d)
+            return (state_sums / counts_view).to(dtype=states.dtype)
+
+        pooled_k_chunks = compute_pooled_states_via_cumsum(k_docs, chunk_counts_view, chunk_lengths)
+        pooled_v_chunks = compute_pooled_states_via_cumsum(v_docs, chunk_counts_view, chunk_lengths)
+        return pooled_k_chunks, pooled_v_chunks
+
+    def sequence_pooling_qkv(self, query_states, key_states, value_states, doc_indices, global_chunk_ids):
+        q_docs = query_states[doc_indices[:, 0], :, doc_indices[:, 1]]
+        k_docs = key_states[doc_indices[:, 0], :, doc_indices[:, 1]]
+        v_docs = value_states[doc_indices[:, 0], :, doc_indices[:, 1]]
+        unique_global_chunk_ids, chunk_lengths = torch.unique_consecutive(global_chunk_ids, return_counts=True)
+        
+        num_unique_chunks = unique_global_chunk_ids.shape[0]
+        chunk_counts_view = chunk_lengths.view(-1, 1, 1).to(dtype=torch.float32)
+        def compute_pooled_states_via_cumsum(states, counts_view, lengths):
+            b, h, d = states.shape
+            states_flat = states.reshape(b, -1).to(dtype=torch.float32)
+            states_cumsum = F.pad(torch.cumsum(states_flat, dim=0), (0, 0, 1, 0))
+            chunk_cu_seqlens = F.pad(torch.cumsum(lengths, 0), (1, 0))
+            
+            state_sums_flat = states_cumsum[chunk_cu_seqlens[1:]] - states_cumsum[chunk_cu_seqlens[:-1]]
+            state_sums = state_sums_flat.view(num_unique_chunks, h, d)
+            return (state_sums / counts_view).to(dtype=states.dtype)
+
+        pooled_q_chunks = compute_pooled_states_via_cumsum(q_docs, chunk_counts_view, chunk_lengths)
+        pooled_k_chunks = compute_pooled_states_via_cumsum(k_docs, chunk_counts_view, chunk_lengths)
+        pooled_v_chunks = compute_pooled_states_via_cumsum(v_docs, chunk_counts_view, chunk_lengths)
+        return pooled_q_chunks, pooled_k_chunks, pooled_v_chunks
+
+    def count_chunks_per_batch(self, doc_ids, attention_mask, kernel_size):
+        batch_size = doc_ids.size(0)
+        chunk_counts = []
+
+        for i in range(batch_size):
+            mask = attention_mask[i]
+            ids = doc_ids[i]
+            valid_ids = ids[mask == 1]
+            
+            if len(valid_ids) == 0:
+                chunk_counts.append(0)
+                continue
+            _, counts = torch.unique_consecutive(valid_ids, return_counts=True)
+            
+            num_chunks = (counts + kernel_size - 1) // kernel_size
+            total_chunks = num_chunks.sum().item()
+            chunk_counts.append(total_chunks)
+
+        return torch.LongTensor(chunk_counts).to(doc_ids.device)
+
+    def _forward(
+        self,
+        hidden_states: torch.Tensor,
+        doc_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        bsz, q_len, _ = hidden_states.shape
+        device, dtype = hidden_states.device, hidden_states.dtype
+        hidden_shape = (bsz, q_len, -1, self.head_dim)
+        
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        routing_query_mask = (doc_ids == 0) & (attention_mask == 1)
+        doc_token_mask = (doc_ids > 0) & (attention_mask == 1)
+
+        query_indices = torch.nonzero(routing_query_mask, as_tuple=False)
+        doc_indices = torch.nonzero(doc_token_mask, as_tuple=False)
+        
+        if doc_indices.shape[0] == 0 or query_indices.shape[0] == 0:
+            raise ValueError("No query or doc tokens found")
+        
+        max_doc_id = int(doc_ids.max().item())
+        attn_output = torch.zeros((bsz, q_len, self.config.num_attention_heads * self.head_dim), device=device, dtype=dtype)
+        if self.is_router_layer:
+            original_doc_ids = doc_ids[doc_token_mask]
+            original_doc_batch_indices = doc_indices[:, 0]
+            
+            global_doc_ids = original_doc_batch_indices * (max_doc_id + 1) + original_doc_ids
+            _, counts = torch.unique_consecutive(global_doc_ids, return_counts=True)
+            total_doc_tokens = global_doc_ids.shape[0]
+            
+            offsets = torch.zeros(counts.shape[0] + 1, dtype=counts.dtype, device=device)
+            offsets[1:] = counts.cumsum(0)
+            offsets = offsets[:-1]
+
+            expanded_offsets = torch.repeat_interleave(offsets, counts)
+            original_order_ranks = torch.arange(total_doc_tokens, device=device) - expanded_offsets
+            
+            chunk_indices = original_order_ranks // self.pooling_kernel_size
+            max_chunks_per_doc = (q_len // self.pooling_kernel_size) + 1
+
+            global_chunk_ids = global_doc_ids * max_chunks_per_doc + chunk_indices
+            
+            routing_q_states = None
+            routing_pooled_k_chunks = None
+
+            if self.decouple_router:
+                routing_q_states = self.router_q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+                if "INFONCE" in self.aux_loss_method:
+                    routing_q_states = F.normalize(routing_q_states, p=2, dim=-1)
+
+
+                r_k_raw = self.router_k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+                r_k_raw = repeat_kv(r_k_raw, self.num_key_value_groups)
+                r_k_docs = r_k_raw[doc_indices[:, 0], :, doc_indices[:, 1]]
+                
+                unique_global_chunk_ids = torch.unique_consecutive(global_chunk_ids)
+                _, chunk_lengths = torch.unique_consecutive(global_chunk_ids, return_counts=True)
+                
+                chunk_counts_view = chunk_lengths.view(-1, 1, 1).to(dtype=torch.float32)
+                b_k, h_k, d_k = r_k_docs.shape
+                k_flat = r_k_docs.reshape(b_k, -1).to(dtype=torch.float32)
+                k_cumsum = F.pad(torch.cumsum(k_flat, dim=0), (0, 0, 1, 0))
+                chunk_cu_seqlens = F.pad(torch.cumsum(chunk_lengths, 0), (1, 0))
+                k_sums_flat = k_cumsum[chunk_cu_seqlens[1:]] - k_cumsum[chunk_cu_seqlens[:-1]]
+                routing_pooled_k_chunks = (k_sums_flat.view(unique_global_chunk_ids.shape[0], h_k, d_k) / chunk_counts_view).to(dtype=r_k_docs.dtype)
+                if "INFONCE" in self.aux_loss_method:
+                    routing_pooled_k_chunks = F.normalize(routing_pooled_k_chunks, p=2, dim=-1)
+
+                pooled_q_chunks = query_states[doc_indices[:, 0], :, doc_indices[:, 1]]
+                pooled_k_chunks = key_states[doc_indices[:, 0], :, doc_indices[:, 1]]
+                pooled_v_chunks = value_states[doc_indices[:, 0], :, doc_indices[:, 1]]
+                num_doc_tokens = pooled_q_chunks.shape[0]
+                num_chunks = num_doc_tokens // self.pooling_kernel_size
+
+                pooled_q_chunks = pooled_q_chunks.view(num_chunks, self.pooling_kernel_size, self.num_heads, self.head_dim).mean(dim=1)
+                pooled_k_chunks = pooled_k_chunks.view(num_chunks, self.pooling_kernel_size, self.num_heads, self.head_dim).mean(dim=1)
+                pooled_v_chunks = pooled_v_chunks.view(num_chunks, self.pooling_kernel_size, self.num_heads, self.head_dim).mean(dim=1)
+                num_heads = self.config.num_attention_heads
+                head_dim = self.head_dim
+            else:
+                pooled_q_chunks, pooled_k_chunks, pooled_v_chunks = self.sequence_pooling_qkv(
+                    query_states, 
+                    key_states, 
+                    value_states, 
+                    doc_indices,
+                    global_chunk_ids,
+                )
+                num_heads = self.config.num_attention_heads
+                head_dim = self.head_dim
+                
+                routing_q_states = query_states
+                routing_pooled_k_chunks = pooled_k_chunks
+                if "INFONCE" in self.aux_loss_method:
+                    routing_q_states = F.normalize(routing_q_states, p=2, dim=-1)
+                    routing_pooled_k_chunks = F.normalize(routing_pooled_k_chunks, p=2, dim=-1)
+                
+            unique_global_chunk_ids = torch.unique_consecutive(global_chunk_ids)
+            num_unique_chunks = unique_global_chunk_ids.shape[0]
+            chunks_per_sample = self.count_chunks_per_batch(doc_ids, doc_token_mask, kernel_size=self.pooling_kernel_size)
+
+            max_chunks = chunks_per_sample.max().item()
+            pooled_router_k_bched = torch.zeros((bsz, max_chunks, num_heads, self.head_dim), device=device, dtype=dtype)
+            chunk_mask = torch.arange(max_chunks, device=device).unsqueeze(0) < chunks_per_sample.unsqueeze(1)
+            pooled_router_k_bched[chunk_mask] = routing_pooled_k_chunks
+            q_lens = routing_query_mask.sum(dim=1) # (B,)
+            max_q_len = int(q_lens.max().item())
+
+            if max_q_len == 0:
+                max_q_len = 1 
+            valid_q_flat = routing_q_states.transpose(1, 2)[routing_query_mask] # [Total_Valid_Q, H, D]
+
+            compact_q_states_t = torch.zeros(
+                bsz, max_q_len, self.config.num_attention_heads, self.head_dim, 
+                device=device, dtype=dtype
+            )
+            
+            idx_range = torch.arange(max_q_len, device=device).unsqueeze(0)
+            mask_compact_q = idx_range < q_lens.unsqueeze(1)
+            
+            compact_q_states_t[mask_compact_q] = valid_q_flat
+            compact_q_states = compact_q_states_t.transpose(1, 2)
+
+            max_scores_per_chunk = self._calculate_routing_scores_adaptive(
+                compact_q_states,         # (B, H, S, D)
+                pooled_router_k_bched,       # (B, C, H, D)
+                mask_compact_q,   # (B, S)
+                chunk_mask            # (B, C)
+            )
+            
+            pooled_global_doc_ids = unique_global_chunk_ids // max_chunks_per_doc
+            pooled_doc_ids_in_sample = pooled_global_doc_ids % (max_doc_id + 1)
+            chunk_to_doc_id_flat = pooled_doc_ids_in_sample # 形状: (total_chunks, )
+
+            chunk_to_doc_id_bched = torch.full((bsz, max_chunks), 0, dtype=torch.long, device=device)
+            chunk_to_doc_id_bched[chunk_mask] = chunk_to_doc_id_flat
+
+            offsets = torch.arange(bsz, device=device) * (max_doc_id + 1)
+            global_chunk_to_doc_id = chunk_to_doc_id_bched + offsets.unsqueeze(1)
+            flat_doc_scores = torch.full((bsz * (max_doc_id + 1),), -float('inf'), device=device, dtype=dtype)
+
+            valid_scores_flat = max_scores_per_chunk[chunk_mask]
+            valid_global_doc_ids_flat = global_chunk_to_doc_id[chunk_mask]
+
+            if self.chunk_reduce_method == "max":
+                doc_scores = flat_doc_scores.scatter_reduce(
+                    dim=0, 
+                    index=valid_global_doc_ids_flat, 
+                    src=valid_scores_flat, 
+                    reduce="amax", 
+                    include_self=True
+                )
+            
+            elif self.chunk_reduce_method == "mean":
+                flat_doc_sums = torch.zeros_like(flat_doc_scores)
+                
+                flat_doc_sums = flat_doc_sums.scatter_reduce(
+                    dim=0,
+                    index=valid_global_doc_ids_flat,
+                    src=valid_scores_flat,
+                    reduce="sum",
+                    include_self=False 
+                )
+                
+                flat_doc_counts = torch.zeros_like(flat_doc_scores)
+                ones = torch.ones_like(valid_scores_flat)
+                
+                flat_doc_counts = flat_doc_counts.scatter_reduce(
+                    dim=0,
+                    index=valid_global_doc_ids_flat,
+                    src=ones,
+                    reduce="sum",
+                    include_self=False
+                )
+                
+                flat_doc_counts_safe = flat_doc_counts.clamp(min=1.0)
+                mean_scores = flat_doc_sums / flat_doc_counts_safe
+                
+                doc_scores = torch.where(
+                    flat_doc_counts > 0,
+                    mean_scores,
+                    flat_doc_scores # 这里是 -inf
+                )
+
+            else:
+                raise ValueError(f"Invalid chunk reduction method: {self.chunk_reduce_method}")
+
+            scores_by_batch = doc_scores.view(bsz, -1)
+            return_scores_by_batch = scores_by_batch.clone()
+
+            num_docs_per_sample = (scores_by_batch > -1e9).sum(dim=1)
+            # 为每个样本计算k值：取配置的top_k和实际文档数的较小者
+            k_per_sample = torch.min(num_docs_per_sample, torch.full_like(num_docs_per_sample, self.top_k_docs))
+
+            _, sorted_indices = torch.sort(scores_by_batch, dim=1, descending=True)
+            
+            range_tensor = torch.arange(scores_by_batch.shape[1], device=device).expand(bsz, -1)
+            selection_mask = range_tensor < k_per_sample.unsqueeze(1)
+            
+            selected_docs_indices = sorted_indices.masked_fill(~selection_mask, -50)
+            
+            prompt_and_response_mask = (doc_ids < 1) & (attention_mask == 1)
+            # 此处的 selected_docs_indices 已经是修复后的张量，所以这行代码无需修改
+            selected_docs_mask = torch.any(doc_ids.unsqueeze(-1) == selected_docs_indices.unsqueeze(1), dim=-1) & doc_token_mask
+
+            pa_indices = torch.nonzero(prompt_and_response_mask, as_tuple=False)
+            q_pa_flat = query_states[pa_indices[:, 0], :, pa_indices[:, 1]]
+            k_pa_flat = key_states[pa_indices[:, 0], :, pa_indices[:, 1]]
+            v_pa_flat = value_states[pa_indices[:, 0], :, pa_indices[:, 1]]
+            sort_key_pa = pa_indices[:, 0] * q_len + pa_indices[:, 1]
+            
+            selected_doc_token_indices = torch.nonzero(selected_docs_mask, as_tuple=False)
+            is_doc_token_mask_flat = doc_token_mask.flatten()
+            global_chunk_ids_padded = torch.full((bsz * q_len,), -1, dtype=torch.long, device=device)
+            global_chunk_ids_padded[is_doc_token_mask_flat] = global_chunk_ids
+            selected_chunk_ids_flat = global_chunk_ids_padded.view(bsz, q_len)[selected_docs_mask]
+            
+            unique_selected_chunk_ids, inverse_indices_fix = torch.unique(selected_chunk_ids_flat, sorted=True, return_inverse=True)
+            if unique_selected_chunk_ids.numel() > 0:
+                first_occurrence_indices = torch.empty_like(unique_selected_chunk_ids, dtype=torch.long)
+                first_occurrence_indices.scatter_reduce_(src=torch.arange(selected_chunk_ids_flat.numel(), device=device),index=inverse_indices_fix, dim=0, reduce='amin', include_self=False)
+                
+                representative_indices = selected_doc_token_indices[first_occurrence_indices]
+                sort_key_chunks = representative_indices[:, 0] * q_len + representative_indices[:, 1]
+
+                map_gcid_to_poolidx = torch.full((int(global_chunk_ids.max().item()) + 1,), -1, dtype=torch.long, device=device)
+                map_gcid_to_poolidx[unique_global_chunk_ids] = torch.arange(num_unique_chunks, device=device)
+                
+                pool_indices_to_gather = map_gcid_to_poolidx[unique_selected_chunk_ids]
+                
+                assert (pool_indices_to_gather.sort().values != pool_indices_to_gather).sum() == 0
+                q_pooled_sel_flat = pooled_q_chunks[pool_indices_to_gather]
+                k_pooled_sel_flat = pooled_k_chunks[pool_indices_to_gather]
+                v_pooled_sel_flat = pooled_v_chunks[pool_indices_to_gather]
+
+                batch_indices_chunks = representative_indices[:, 0]
+            else:
+                sort_key_chunks = torch.tensor([], dtype=torch.long, device=device)
+                q_pooled_sel_flat = torch.tensor([], dtype=dtype, device=device).view(0, num_heads, head_dim)
+                k_pooled_sel_flat = torch.tensor([], dtype=dtype, device=device).view(0, num_heads, head_dim)
+                v_pooled_sel_flat = torch.tensor([], dtype=dtype, device=device).view(0, num_heads, head_dim)
+                batch_indices_chunks = torch.tensor([], dtype=torch.long, device=device)
+
+            q_combined = torch.cat([q_pa_flat, q_pooled_sel_flat], dim=0)
+            k_combined = torch.cat([k_pa_flat, k_pooled_sel_flat], dim=0)
+            v_combined = torch.cat([v_pa_flat, v_pooled_sel_flat], dim=0)
+            
+            combined_sort_keys = torch.cat([sort_key_pa, sort_key_chunks], dim=0)
+            _, final_sort_indices = torch.sort(combined_sort_keys)
+            
+            q_a_final = q_combined[final_sort_indices]
+            k_a_final = k_combined[final_sort_indices]
+            v_a_final = v_combined[final_sort_indices]
+            
+            # 4.4 计算cu_seqlens (逻辑不变)
+            batch_indices_pa = pa_indices[:, 0]
+            batch_indices_combined = torch.cat([batch_indices_pa, batch_indices_chunks], dim=0)
+            sorted_batch_indices = batch_indices_combined[final_sort_indices]
+            
+            batch_counts_a = torch.bincount(sorted_batch_indices, minlength=bsz)
+            cu_seqlens_a = F.pad(torch.cumsum(batch_counts_a, dim=0, dtype=torch.int32), (1, 0))
+        else:
+            prompt_and_response_mask = (doc_ids < 1) & (attention_mask == 1)
+            pa_indices = torch.nonzero(prompt_and_response_mask, as_tuple=False)
+            q_a_final = query_states[pa_indices[:, 0], :, pa_indices[:, 1]]
+            k_a_final = key_states[pa_indices[:, 0], :, pa_indices[:, 1]]
+            v_a_final = value_states[pa_indices[:, 0], :, pa_indices[:, 1]]
+            batch_counts_a = prompt_and_response_mask.sum(dim=1)
+            cu_seqlens_a = F.pad(torch.cumsum(batch_counts_a, dim=0, dtype=torch.int32), (1, 0))
+            return_scores_by_batch = None
+        
+        if q_a_final.shape[0] > 0:
+            output_a_final = flash_attn_varlen_func(
+                q_a_final, k_a_final, v_a_final,
+                cu_seqlens_q=cu_seqlens_a, cu_seqlens_k=cu_seqlens_a,
+                max_seqlen_q=int(batch_counts_a.max()), max_seqlen_k=int(batch_counts_a.max()),
+                dropout_p=self.attention_dropout if self.training else 0.0,
+                causal=True
+            ).view(-1, self.config.num_attention_heads * self.head_dim)
+            
+            if self.is_router_layer:
+                is_pa_mask_combined = torch.cat([
+                    torch.ones(pa_indices.shape[0], dtype=torch.bool, device=device),
+                    torch.zeros(q_pooled_sel_flat.shape[0], dtype=torch.bool, device=device) # 修正为使用池化块的数量
+                ], dim=0)
+                is_pa_mask_sorted = is_pa_mask_combined[final_sort_indices]
+                
+                output_pa_part = output_a_final[is_pa_mask_sorted]
+                attn_output[pa_indices[:, 0], pa_indices[:, 1]] = output_pa_part
+            else:
+                attn_output[pa_indices[:, 0], pa_indices[:, 1]] = output_a_final
+
+        indices_b = torch.nonzero(doc_token_mask, as_tuple=False)
+        if indices_b.shape[0] > 0:
+            q_b, k_b, v_b = query_states[indices_b[:, 0], :, indices_b[:, 1]], key_states[indices_b[:, 0], :, indices_b[:, 1]], value_states[indices_b[:, 0], :, indices_b[:, 1]]
+            doc_ids_b = doc_ids[indices_b[:, 0], indices_b[:, 1]]
+            batch_indices_b = indices_b[:, 0]
+            
+            global_doc_ids_b = batch_indices_b * (max_doc_id + 1) + doc_ids_b
+            
+            _, counts_b = torch.unique_consecutive(global_doc_ids_b, return_counts=True)
+            cu_seqlens_b = F.pad(torch.cumsum(counts_b, dim=0, dtype=torch.int32), (1, 0))
+            
+            output_b_flat = flash_attn_varlen_func(
+                q_b, k_b, v_b, cu_seqlens_q=cu_seqlens_b, cu_seqlens_k=cu_seqlens_b,
+                max_seqlen_q=int(counts_b.max()), max_seqlen_k=int(counts_b.max()),
+                dropout_p=self.attention_dropout if self.training else 0.0, causal=True
+            ).view(-1, self.config.num_attention_heads * self.head_dim)
+            
+            attn_output[indices_b[:, 0], indices_b[:, 1]] += output_b_flat
+
+        return (self.o_proj(attn_output), return_scores_by_batch), None

src/msa/model.py

@@ -0,0 +1,733 @@
+import sys
+import os
+sys.path.append(os.getcwd())
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Union, Optional, Tuple, List, Dict
+from functools import partial
+from dataclasses import dataclass
+from transformers.modeling_outputs import ModelOutput
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3Model, 
+    FlashAttentionKwargs,
+    Qwen3Config,
+    Qwen3RMSNorm,
+    Qwen3RotaryEmbedding,
+    DynamicCache,
+    Qwen3DecoderLayer,
+    Qwen3MLP,
+    Qwen3Attention,
+    Qwen3PreTrainedModel,
+)
+from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast
+from transformers.processing_utils import Unpack
+from transformers.cache_utils import Cache
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+from src.msa import MemorySparseAttention, MSAGenerationMixin, MSAConfig
+
+@dataclass
+class MSALayerModelOutputWithPast(ModelOutput):
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+    all_docs_scores: Optional[Dict] = None
+
+@dataclass
+class MSACausalLMOutputWithPast(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    lm_loss: Optional[torch.FloatTensor] = None
+    aux_loss: Optional[torch.FloatTensor] = None
+    answer_loss: Optional[torch.FloatTensor] = None
+    reconstruction_loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+    temperature: Optional[torch.FloatTensor] = None
+    train_router_metrics: Optional[Dict] = None
+
+class MSADeocoderLayer(Qwen3DecoderLayer):
+    def __init__(self, config: Qwen3Config, layer_idx: int, attn_type: str = "sparse_attention"):
+        super().__init__(config=config, layer_idx=layer_idx)
+        self.layer_idx = layer_idx
+        self.attn_type = attn_type
+        self.hidden_size = config.hidden_size
+        if attn_type == "full_attention":
+            self.self_attn = Qwen3Attention(config=config, layer_idx=layer_idx)
+        elif attn_type == "sparse_attention":
+            self.self_attn = MemorySparseAttention(config=config, layer_idx=layer_idx)
+        
+        self.mlp = Qwen3MLP(config)
+        self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        config.sliding_window = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        output_docs_score: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
+        doc_ids: Optional[torch.Tensor] = None,
+        input_ids: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        residual = hidden_states
+        
+        hidden_states = self.input_layernorm(hidden_states)
+        
+        # Self Attention
+        if self.attn_type == "full_attention":
+            hidden_states, self_attn_weights = self.self_attn(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                doc_ids=doc_ids,
+                input_ids=input_ids,
+                **kwargs,
+            )
+        else:
+            hidden_states, self_attn_weights = self.self_attn(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                doc_ids=doc_ids,
+                input_ids=input_ids,
+                **kwargs,
+            )
+        
+        if isinstance(hidden_states, tuple):
+            hidden_states, docs_score = hidden_states
+        else:
+            docs_score = None
+        hidden_states = residual + hidden_states
+        
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+        if output_attentions:
+            outputs += (self_attn_weights,)
+        
+        if output_docs_score:
+            outputs += (docs_score,)
+
+        return outputs
+
+class MSAModel(Qwen3Model):
+    def __init__(self, config: Qwen3Config):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.rewrite_position = config.msa_config.rewrite_position
+        self.layers = nn.ModuleList([
+            MSADeocoderLayer(config, layer_idx, attn_type="sparse_attention") 
+            for layer_idx in range(config.num_hidden_layers)
+        ])
+        self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen3RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+    
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_docs_score: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        doc_ids: Optional[torch.LongTensor] = None,
+        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+    ) -> BaseModelOutputWithPast:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            use_cache = False
+
+        # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache
+        if not isinstance(past_key_values, (type(None), Cache)):
+            raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+        
+        if not self.rewrite_position and self.training:
+            position_ids = None
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # causal_mask = self._update_causal_mask(
+        #     attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+        # )
+        causal_mask = attention_mask
+
+        hidden_states = inputs_embeds
+        # import pdb;pdb.set_trace()
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        all_docs_scores = () if output_docs_score else None
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    partial(decoder_layer.__call__, **flash_attn_kwargs),
+                    hidden_states,
+                    causal_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    output_docs_score,
+                    use_cache,
+                    cache_position,
+                    position_embeddings,
+                    doc_ids,
+                    input_ids,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=causal_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    output_docs_score=output_docs_score,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                    position_embeddings=position_embeddings,
+                    doc_ids=doc_ids,
+                    input_ids=input_ids,
+                    **flash_attn_kwargs,
+                )
+            
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+            if output_docs_score:
+                all_docs_scores += (layer_outputs[-1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        return MSALayerModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+            all_docs_scores=all_docs_scores
+        )
+    
+class MSAForCausalLM(Qwen3PreTrainedModel, MSAGenerationMixin):
+    config_class = MSAConfig
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_layers = config.num_hidden_layers
+        self.router_layer_idx = config.msa_config.router_layer_idx
+
+        if self.router_layer_idx == "all":
+            self.router_layer_idx = list(range(config.num_hidden_layers))
+        else:
+            self.router_layer_idx = [int(i) for i in self.router_layer_idx.split(",")]
+
+        self.mid_layers = config.num_hidden_layers // 2
+        self.model = MSAModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.aux_loss = config.msa_config.aux_loss
+        self.lmloss_weigth = config.msa_config.lmloss_weigth
+        self.auxloss_weight = config.msa_config.auxloss_weight
+        self.recloss_weight = config.msa_config.recloss_weight
+        self.ansloss_weight = config.msa_config.ansloss_weight
+        self.aux_loss_method = config.msa_config.aux_loss_method  # INFONCE, BCE, INFONCE_DECOUPLE, INFONCE_DECOUPLE_FOCAL 
+        self.decouple_router = config.msa_config.decouple_router
+            
+        if "INFONCE" in self.aux_loss_method:
+            temperature = config.msa_config.infonce_loss_temp
+            self.temperature = nn.Parameter(torch.ones([]) * temperature, requires_grad=False)
+        elif self.aux_loss_method == "BCE":
+            self.b = nn.Parameter(-20 * torch.ones([]), requires_grad=True)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_docs_score: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        # msa
+        doc_ids: Optional[torch.LongTensor] = None,
+        batch_aux_labels: List[List[int]] = None,
+        batch_reconstruction_labels: Optional[torch.LongTensor] = None,
+        batch_answer_labels: Optional[torch.LongTensor] = None,
+        train_qa_samples: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        r"""
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+            logits_to_keep (`int` or `torch.Tensor`, *optional*):
+                If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+                `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+                token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+                If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+                This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, Qwen3ForCausalLM
+
+        >>> model = Qwen3ForCausalLM.from_pretrained("Qwen/Qwen3-8B")
+        >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        
+        output_docs_score = self.aux_loss
+        
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_docs_score=output_docs_score,
+            cache_position=cache_position,
+            doc_ids=doc_ids,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+        
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        kept_hidden_states = hidden_states[:, slice_indices, :]
+
+        shift_labels = kwargs.pop("shift_labels", None)
+        logits = None
+        loss = None
+        reconstruction_loss = None
+        aux_loss = None
+        answer_loss = None
+        # if in training mode, don't materialize logits
+        if self.training and (labels is not None or shift_labels is not None):
+            loss = LigerForCausalLMLoss(
+                hidden_states=kept_hidden_states,
+                lm_head_weight=self.lm_head.weight,
+                labels=labels,
+                shift_labels=shift_labels,
+                hidden_size=self.config.hidden_size,
+                **kwargs,
+            )
+            if batch_reconstruction_labels is not None:
+                reconstruction_loss = LigerForCausalLMLoss(
+                    hidden_states=kept_hidden_states,
+                    lm_head_weight=self.lm_head.weight,
+                    labels=batch_reconstruction_labels,
+                    shift_labels=None,
+                    hidden_size=self.config.hidden_size,
+                    **kwargs,
+                )
+            else:
+                reconstruction_loss = torch.tensor(0.0).to(hidden_states.device)
+
+            if torch.sum(train_qa_samples) == 0:
+                batch_answer_labels = None
+            else:
+                train_qa_samples_mask = train_qa_samples == 1
+                temp_kept_hidden_states = kept_hidden_states[train_qa_samples_mask]
+                temp_batch_answer_labels = batch_answer_labels[train_qa_samples_mask]
+
+            if batch_answer_labels is not None:
+                answer_loss = LigerForCausalLMLoss(
+                    hidden_states=temp_kept_hidden_states,
+                    lm_head_weight=self.lm_head.weight,
+                    labels=temp_batch_answer_labels,
+                    shift_labels=None,
+                    hidden_size=self.config.hidden_size,
+                    **kwargs,
+                )
+            else:
+                answer_loss = torch.tensor(0.0).to(hidden_states.device)
+            
+
+        else:  # if in inference mode materialize logits
+            logits = self.lm_head(kept_hidden_states)
+            if labels is not None:
+                loss = self.loss_function(
+                    logits=logits,
+                    labels=labels,
+                    vocab_size=self.config.vocab_size,
+                    **kwargs,
+                )
+                if batch_reconstruction_labels is not None:
+                    reconstruction_loss = self.loss_function(
+                        logits=logits,
+                        labels=batch_reconstruction_labels,
+                        vocab_size=self.config.vocab_size,
+                        **kwargs,
+                    )
+                else:
+                    reconstruction_loss = torch.tensor(0.0).to(hidden_states.device)
+
+                if batch_answer_labels is not None:
+                    answer_loss = self.loss_function(
+                        logits=logits,
+                        labels=batch_answer_labels,
+                        vocab_size=self.config.vocab_size,
+                        **kwargs,
+                    )
+                else:
+                    answer_loss = torch.tensor(0.0).to(hidden_states.device)
+            
+        lm_loss = torch.tensor(0.0).to(hidden_states.device)
+        if loss is not None:
+            lm_loss = loss.clone()
+        aux_loss = torch.tensor(0.0).to(hidden_states.device)
+        train_router_metrics = None
+        if batch_aux_labels is not None and self.aux_loss:
+            aux_loss, train_router_metrics = self.caculate_aux_loss(aux_loss, outputs, batch_aux_labels, hidden_states.device, hidden_states.dtype)
+
+        # Ensure all loss components have default values
+        reconstruction_loss = reconstruction_loss if reconstruction_loss is not None else torch.tensor(0.0).to(hidden_states.device)
+        answer_loss = answer_loss if answer_loss is not None else torch.tensor(0.0).to(hidden_states.device)
+        aux_loss = aux_loss if aux_loss is not None else torch.tensor(0.0).to(hidden_states.device)
+        
+        if loss is not None:
+            loss = self.lmloss_weigth * loss + \
+                   self.recloss_weight * reconstruction_loss + \
+                   self.auxloss_weight * aux_loss + \
+                   self.ansloss_weight * answer_loss
+                   
+        return MSACausalLMOutputWithPast(
+            loss=loss,
+            lm_loss=lm_loss,
+            aux_loss=aux_loss,
+            answer_loss=answer_loss,
+            reconstruction_loss=reconstruction_loss,
+            train_router_metrics=train_router_metrics,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            temperature=self.temperature if "INFONCE" in self.aux_loss_method else torch.tensor(0.0).to(hidden_states.device),
+        )
+
+    def calculate_decoupled_infonce_loss(self, logits, label, num_pos):
+        """
+        改进版：Decoupled InfoNCE
+        解决了多正样本之间的互斥问题，检索任务推荐使用。
+        """
+        with torch.no_grad():
+            temperature = self.temperature.clamp(0.001, 0.5)
+        
+        # 1. 缩放
+        scaled_logits = logits / temperature
+        
+        # 2. 数值稳定：转为 exp 域
+        # 减去最大值防止溢出 (标准 Softmax trick)
+        max_logits = torch.max(scaled_logits, dim=0, keepdim=True)[0].detach()
+        exp_logits = torch.exp(scaled_logits - max_logits)
+        
+        # 获取所有负样本的 exp 之和
+        neg_exp_sum = torch.sum(exp_logits * (1 - label), dim=0, keepdim=True)
+        
+        # 4. 计算 Log Prob
+        # 对于每一个正样本 i:
+        # Prob_i = exp_i / (exp_i + neg_exp_sum)
+        # 这样分母里就没有“其他正样本”在竞争了
+        denominators = exp_logits + neg_exp_sum
+        
+        log_probs = scaled_logits - max_logits - torch.log(denominators + 1e-10)
+        
+        # 5. 计算 Loss
+        # 只取正样本位置的 log_prob
+        loss_map = - log_probs * (label / num_pos.clamp(min=1.0))
+        
+        return loss_map.sum(dim=0)
+
+    def caculate_infonce_loss(self, logits, label, num_pos):
+        # 限制温度范围避免数值不稳定
+        with torch.no_grad():
+            temperature = self.temperature.clamp(0.001,0.5)
+        # 对logits进行温度缩放
+        scaled_logits = logits / temperature
+        # 确保正样本数至少为1，避免除零错误
+        safe_num_pos = num_pos.clamp(min=1.0) 
+        # 将标签转换为与缩放后logits相同的浮点类型
+        aux_label_float = label.to(dtype=scaled_logits.dtype)
+        # 计算InfoNCE损失：负的对数softmax概率与标签的加权和
+        one_aux_loss = -torch.sum(F.log_softmax(scaled_logits, dim=0) * (aux_label_float / safe_num_pos), dim=0)
+        return one_aux_loss
+
+    def caculate_bce_loss(self, logits, label):
+        label[label==0] = -1 # 将 0 (负样本) 映射为 -1
+        one_aux_loss = -torch.mean(F.logsigmoid((logits + self.b) * label))
+        return one_aux_loss
+
+    def calculate_multi_pos_focal_infonce(self, logits, label, gamma=2.0):
+        """
+        Args:
+            logits: (Batch_Size, Num_Candidates) 或者是 (N, 1) 的形式
+            label: (Batch_Size, Num_Candidates) Multi-hot 标签，1为正，0为负
+            gamma: Focal 参数
+        """
+        with torch.no_grad():
+            temperature = self.temperature.clamp(0.001, 0.5)
+        
+        # 1. 缩放 logits
+        scaled_logits = logits / temperature
+        
+        # 2. 核心 Trick：数值稳定地将 logits 转为 exp 域
+        # 为了防止 exp 溢出，先减去最大值
+        max_logits = torch.max(scaled_logits, dim=0, keepdim=True)[0].detach()
+        exp_logits = torch.exp(scaled_logits - max_logits)
+        
+        # 3. 构建分母 (Denominator)
+        # 传统的 InfoNCE 分母是 sum(exp_logits)
+        # 我们现在的分母应该是：当前正样本的 exp + 所有负样本的 exp
+        # 等价于：总和 - 其他正样本的 exp
+        
+        sum_exp = torch.sum(exp_logits, dim=0, keepdim=True) # 所有样本的 exp 之和
+        
+        # 利用 label (multi-hot) 找出所有正样本的 exp
+        # label 为 1 的位置是正样本
+        pos_exp = exp_logits * label 
+        neg_exp_sum = torch.sum(exp_logits * (1 - label), dim=0, keepdim=True) # 所有负样本的 exp 之和
+        
+        # 4. 计算每个正样本对应的 Softmax 概率
+        # 对于每一个位置 i (如果是正样本)：
+        # Prob_i = exp_i / (exp_i + sum(exp_negatives))
+        # 注意：这里分母不包含 label 中其他的正样本！
+        
+        # 这里利用广播机制：
+        # 分母 = 当前位置的 exp (如果是正样本) + 所有负样本的 sum_exp
+        denominators = exp_logits + neg_exp_sum
+        
+        # 计算概率 P
+        probs = exp_logits / denominators
+        
+        # 5. 计算 Log Softmax (数值更稳定)
+        # log(p) = logits - log(denominators)
+        # 同样使用数值稳定的 logaddexp 或者直接操作
+        # 这里为了代码清晰，直接对上面算出的 probs 取 log，实际工程中建议用 log_softmax 形式优化
+        log_probs = torch.log(probs + 1e-10)
+        
+        # 6. Focal Weight
+        # w = (1 - p)^gamma
+        focal_weights = (1 - probs).pow(gamma)
+        
+        # 7. 计算 Loss
+        # 只保留正样本位置的 Loss
+        loss_map = - focal_weights * log_probs * label
+        
+        # 8. 归一化
+        # 除以正样本的总数
+        num_pos = label.sum()
+        loss = loss_map.sum() / (num_pos + 1e-6)
+        return loss
+
+    def calculate_focal_infonce_loss(self, logits, label, num_pos, gamma=2.0):
+        """
+        Args:
+            logits: 模型输出的 logits
+            label: 正样本的 mask (通常是 multi-hot)
+            num_pos: 正样本的数量
+            gamma: Focal Loss 的超参数，控制挖掘难样本的程度，通常设为 2.0
+        """
+        with torch.no_grad():
+            # 限制温度范围，防止数值不稳定
+            temperature = self.temperature.clamp(0.001, 0.5)
+        
+        # 1. 缩放 Logits
+        scaled_logits = logits / temperature
+        
+        # 2. 提前计算 Softmax 分数 (即公式中的 P = exp(s)/Z)
+        # dim=0 对应你的输入是单个向量的情况
+        probs = F.softmax(scaled_logits, dim=0)
+        
+        # 3. 计算 Log Softmax (用于 Loss 计算，比 log(probs) 数值更稳定)
+        log_probs = F.log_softmax(scaled_logits, dim=0)
+        
+        # 4. 计算 Focal Weight: w = (1 - p)^gamma
+        # 公式对应：w_i = (1 - p_i)^gamma
+        # 这里不对权重进行 detach，允许梯度回传以调整对难易样本的关注度 (参考原始 Focal Loss 论文)
+        focal_weights = (1 - probs).pow(gamma)
+        
+        # 5. 准备 Label 和归一化项
+        safe_num_pos = num_pos.clamp(min=1.0) 
+        aux_label_float = label.to(dtype=scaled_logits.dtype)
+        
+        # 6. 计算加权的 InfoNCE Loss
+        # 公式对应：Sum [ -w_i * log(p_i) ]
+        # 这里利用 label (0/1) 来只保留正样本的 Loss，并除以 num_pos 做平均
+        weighted_loss = -torch.sum(
+            focal_weights * log_probs * (aux_label_float / safe_num_pos), 
+            dim=0
+        )
+
+        return weighted_loss
+
+    def caculate_aux_loss(self, aux_loss, outputs, batch_aux_labels, device, dtype):
+        count = 0
+        all_layer_doc_score = outputs.all_docs_scores
+        train_router_metrics = {}
+        for layer_idx, layer_doc_score in enumerate(all_layer_doc_score):
+            # if self.mid_layers <= layer_idx < self.num_layers:
+            if self.mid_layers <= layer_idx and layer_idx in self.router_layer_idx:
+                for b in range(len(batch_aux_labels)):
+                    aux_logits_full = layer_doc_score[b] 
+                    aux_label = torch.LongTensor(batch_aux_labels[b]).type(dtype).to(device)
+
+                    # 1. 创建一个 mask 来找到所有有效的 (非 -inf) doc
+                    valid_doc_mask = (aux_logits_full > -1e9) # 或 != -float('inf')
+
+                    # 2. 使用 *同一个 mask* 来过滤 logits 和 labels
+                    aux_logits = aux_logits_full[valid_doc_mask]
+                    
+                    if aux_label.shape[0] == 0:
+                        continue
+                    
+                    num_pos = aux_label.sum()
+                    if self.aux_loss_method == "BCE":
+                        one_aux_loss = self.caculate_bce_loss(aux_logits, aux_label)
+                    elif self.aux_loss_method == "INFONCE":
+                        one_aux_loss = self.caculate_infonce_loss(aux_logits, aux_label, num_pos)
+                    elif self.aux_loss_method == "INFONCE_FOCAL":
+                        one_aux_loss = self.calculate_focal_infonce_loss(aux_logits, aux_label, num_pos)
+                    elif self.aux_loss_method == "INFONCE_DECOUPLE":
+                        one_aux_loss = self.calculate_decoupled_infonce_loss(aux_logits, aux_label, num_pos)
+                    elif self.aux_loss_method == "INFONCE_DECOUPLE_FOCAL":
+                        one_aux_loss = self.calculate_multi_pos_focal_infonce(aux_logits, aux_label, num_pos)
+
+                    aux_loss += one_aux_loss
+                    count += 1
+
+                    # 计算recall
+                    recall_at_n = [1, 5, 10]
+                    for at_n in recall_at_n:
+                        at_n = min(at_n, aux_logits.shape[0])
+                        top_k_indices = torch.topk(aux_logits, at_n)[1].cpu().tolist()
+                        hit_count = 0
+                        for idx in top_k_indices:
+                            if aux_label[idx] == 1:
+                                hit_count += 1
+                        recall_n = hit_count / min(num_pos.item(), 10)
+                        if f'recall@{at_n}' not in train_router_metrics:
+                            train_router_metrics[f'recall@{at_n}'] = [recall_n]
+                        else:
+                            train_router_metrics[f'recall@{at_n}'].append(recall_n)
+        
+        if count != 0:
+            aux_loss = aux_loss / count
+        for k, v in train_router_metrics.items():
+            train_router_metrics[k] = sum(v) / len(v)
+        return aux_loss, train_router_metrics

src/msa_service.py

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

src/prefill.py

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

src/types.py

@@ -0,0 +1,21 @@
+from dataclasses import dataclass
+import multiprocessing as mp
+
+
+@dataclass
+class Document:
+    doc: str = ""
+    doc_id: int = 0
+    num_chunks: int = 0
+
+class ProtocolConstants:
+
+    @staticmethod
+    def expect(q: mp.Queue, constant):
+        k, v = q.get()
+        assert k == constant, f"expect {constant} but got {k}"
+        return v
+    
+    @staticmethod
+    def send(q: mp.Queue, constant, data=None, block=True):
+        q.put((constant, data), block=block)

src/utils/__init__.py

@@ -0,0 +1,15 @@
+from .common import (
+    patch_transformer_logging,
+    print_trainable_params,
+    draw_bounding_boxes,
+    post_process_generate_ids,
+    decode_generate_ids,
+    smart_tokenizer_and_embedding_resize,
+)
+from .callbacks import (
+    ModeltimeCallback,
+    ProfCallback,
+    SacredCallback,
+    ModelEvalCallback,
+    DSEmptyCacheCallback
+)

src/utils/cache.py

@@ -0,0 +1,512 @@
+#!/usr/bin/env python3
+"""
+缓存模块
+包含CustomDynamicCache和CustomQuantizeDynamicCache两个类
+从eval_anything_v2_batch.py中提取出来，实现模块化管理
+"""
+
+import copy
+from typing import Optional, Tuple
+import torch
+from transformers.cache_utils import DynamicCache, QuantoQuantizedCache, QuantizedCacheConfig
+
+
+class CustomDynamicCache(DynamicCache):
+    """
+    自定义动态缓存类
+    扩展标准的DynamicCache，添加额外的元数据存储和查询功能
+    """
+    def __init__(self, _distributed_cache_data=None):
+        super().__init__(_distributed_cache_data)
+        self.cache_kwargs = {}
+        self.group_cache = {}
+        self.meta = {}
+        self.router_key_cache = []
+
+    def clear_kvcache(self):
+        self.key_cache = []
+        self.value_cache = []
+
+    def record_kwargs(self, layer_idx, kwargs):
+        """
+        记录层的元数据信息
+
+        Args:
+            layer_idx: 层索引
+            kwargs: 包含路由层信息的字典
+        """
+        if layer_idx in self.cache_kwargs:
+            self.cache_kwargs[layer_idx].update(kwargs)
+        else:
+            self.cache_kwargs[layer_idx] = kwargs
+
+    def get_layer_length(self):
+        return len(self.cache_kwargs)
+
+    def get_kvcache(self, layer_idx):
+        """
+        获取指定层的KV缓存
+
+        Args:
+            layer_idx: 层索引
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: (key_cache, value_cache)
+        """
+        key_cache = self.key_cache[layer_idx]
+        value_cache = self.value_cache[layer_idx]
+        return key_cache, value_cache
+
+    def get_router_kcache(self, layer_idx):
+        if layer_idx < len(self.router_key_cache):
+            return self.router_key_cache[layer_idx]
+        else:
+            return None
+
+    def clear_query(self):
+        """
+        清理查询相关的临时数据
+        移除查询过程中产生的临时数据，保持缓存清洁
+        """
+        for k, v in self.cache_kwargs.items():
+            if "compacked_key_cache" in v:
+                v.pop("compacked_key_cache")
+                v.pop("compacked_value_cache")
+            if "prefill_stage2_kvcache_size" in v:
+                v.pop("prefill_stage2_kvcache_size")
+            if "prefill_stage1_kvcache_size" in v:
+                v.pop("prefill_stage1_kvcache_size")
+            if "recall_topk" in v:
+                v.pop("recall_topk")
+        return self
+
+    def get_seq_length(self, layer_idx=0) -> int:
+        """
+        返回缓存状态的序列长度
+
+        Args:
+            layer_idx: 可选的层索引
+
+        Returns:
+            int: 序列长度
+        """
+        is_empty_layer = (
+            len(self.key_cache) == 0  # no cache in any layer
+            or len(self.key_cache) <= layer_idx  # skipped `layer_idx` and hasn't run a layer with cache after it
+            or not self.key_cache[layer_idx].numel()  # the layer has no cache
+        )
+        layer_seq_length = self.key_cache[layer_idx].shape[-2] if not is_empty_layer else 0
+        return layer_seq_length
+
+    def copy(self):
+        """
+        创建缓存的深拷贝
+
+        Returns:
+            CustomDynamicCache: 缓存的新副本
+        """
+        new_cache = CustomDynamicCache()
+        new_cache.key_cache = [k.clone() for k in self.key_cache]
+        new_cache.value_cache = [v.clone() for v in self.value_cache]
+        new_cache.cache_kwargs = copy.deepcopy(self.cache_kwargs)
+        new_cache.group_cache = copy.deepcopy(self.group_cache)
+        new_cache.meta = copy.deepcopy(self.meta)
+        new_cache._seen_tokens = self._seen_tokens
+        return new_cache
+
+    def update_router_kcache(
+        self,
+        key_states: torch.Tensor,
+        layer_idx: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        # Update the cache
+        if key_states is not None:
+            if len(self.router_key_cache) <= layer_idx:
+                # There may be skipped layers, fill them with empty lists
+                for _ in range(len(self.router_key_cache), layer_idx):
+                    self.router_key_cache.append(torch.tensor([]))
+                self.router_key_cache.append(key_states)
+            elif (
+                not self.router_key_cache[layer_idx].numel()  # prefers not t.numel() to len(t) == 0 to export the model
+            ):  # fills previously skipped layers; checking for tensor causes errors
+                self.router_key_cache[layer_idx] = key_states
+            else:
+                self.router_key_cache[layer_idx] = torch.cat([self.router_key_cache[layer_idx], key_states], dim=-2)
+
+        return self.router_key_cache[layer_idx]
+    
+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, d)
+
+    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)
+
+    def update_router_kcache(
+        self,
+        key_states: torch.Tensor,
+        layer_idx: int,
+    ) -> 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()
+        return super().update_router_kcache(key_states, layer_idx)
+
+class CustomQuantizeDynamicCache(QuantoQuantizedCache):
+    """
+    自定义量化动态缓存类
+    扩展标准的QuantoQuantizedCache，添加额外的元数据存储和查询功能
+    """
+    def __init__(self, cache_config):
+        super().__init__(cache_config)
+        self.cache_config = cache_config
+        self.cache_kwargs = {}
+        self.group_cache = {}
+        self.meta = {}
+
+    def record_kwargs(self, layer_idx, kwargs):
+        """
+        记录层的元数据信息
+
+        Args:
+            layer_idx: 层索引
+            kwargs: 包含路由层信息的字典
+        """
+        if layer_idx in self.cache_kwargs:
+            self.cache_kwargs[layer_idx].update(kwargs)
+        else:
+            self.cache_kwargs[layer_idx] = kwargs
+
+    def get_layer_length(self):
+        return len(self.cache_kwargs)
+
+    def clear_kvcache(self):
+        self._quantized_key_cache = []
+        self._quantized_value_cache = []
+        self.key_cache = []
+        self.value_cache = []
+
+    def get_kvcache(self, layer_idx):
+        """
+        获取指定层的KV缓存（反量化后）
+
+        Args:
+            layer_idx: 层索引
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: (key_cache, value_cache)
+        """
+        dequant_key = self._dequantize(self._quantized_key_cache[layer_idx])
+        dequant_value = self._dequantize(self._quantized_value_cache[layer_idx])
+        return dequant_key, dequant_value
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs=None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        更新缓存
+
+        Args:
+            key_states: 新的key状态
+            value_states: 新的value状态
+            layer_idx: 层索引
+            cache_kwargs: 缓存关键字参数
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: 更新后的key和value状态
+        """
+        # Update the number of seen tokens
+        if layer_idx == 0:
+            self._seen_tokens += key_states.shape[-2]
+
+        if len(self.key_cache) < layer_idx:
+            for i in range(len(self.key_cache), layer_idx):
+                self.key_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device))
+                self.value_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device))
+                self._quantized_key_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device))
+                self._quantized_value_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device))
+
+        if len(self.key_cache) == layer_idx:
+            self._quantized_key_cache.append(self._quantize(key_states.contiguous(), axis=self.axis_key))
+            self._quantized_value_cache.append(self._quantize(value_states.contiguous(), axis=self.axis_value))
+            self.key_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device))
+            self.value_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device))
+            keys_to_return, values_to_return = key_states, value_states
+        else:
+            dequant_key = self._dequantize(self._quantized_key_cache[layer_idx])
+            dequant_value = self._dequantize(self._quantized_value_cache[layer_idx])
+            keys_to_return = [dequant_key, self.key_cache[layer_idx], key_states]
+            values_to_return = [dequant_value, self.value_cache[layer_idx], value_states]
+
+            keys_to_return = torch.cat(keys_to_return, dim=-2)
+            values_to_return = torch.cat(values_to_return, dim=-2)
+            if (
+                self.key_cache[layer_idx].dim() == 4
+                and self.key_cache[layer_idx].shape[-2] + 1 >= self.residual_length
+            ):
+                self._quantized_key_cache[layer_idx] = self._quantize(keys_to_return.contiguous(), axis=self.axis_key)
+                self._quantized_value_cache[layer_idx] = self._quantize(
+                    values_to_return.contiguous(), axis=self.axis_value
+                )
+                self.key_cache[layer_idx] = torch.zeros(0, dtype=key_states.dtype, device=key_states.device)
+                self.value_cache[layer_idx] = torch.zeros(0, dtype=key_states.dtype, device=key_states.device)
+            else:
+                self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
+                self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)
+
+        return keys_to_return, values_to_return
+
+    def get_seq_length(self, layer_idx=0) -> int:
+        """
+        返回缓存状态的序列长度
+
+        Args:
+            layer_idx: 可选的层索引
+
+        Returns:
+            int: 序列长度
+        """
+        is_empty_layer = (
+            len(self._quantized_key_cache) == 0  # no cache in any layer
+            or len(self._quantized_key_cache) <= layer_idx  # skipped `layer_idx` and hasn't run a layer with cache after it
+            or not self._quantized_key_cache[layer_idx].numel()  # the layer has no cache
+        )
+        layer_seq_length = self._quantized_key_cache[layer_idx].shape[-2] if not is_empty_layer else 0
+        return layer_seq_length
+
+    def clear_query(self):
+        """
+        清理查询相关的临时数据
+        移除查询过程中产生的临时数据，保持缓存清洁
+        """
+        for k, v in self.cache_kwargs.items():
+            if "compacked_key_cache" in v:
+                v.pop("compacked_key_cache")
+                v.pop("compacked_value_cache")
+            if "prefill_stage2_kvcache_size" in v:
+                v.pop("prefill_stage2_kvcache_size")
+            if "prefill_stage1_kvcache_size" in v:
+                v.pop("prefill_stage1_kvcache_size")
+            if "recall_topk" in v:
+                v.pop("recall_topk")
+        return self
+
+    def copy(self):
+        """
+        创建缓存的深拷贝
+
+        Returns:
+            CustomQuantizeDynamicCache: 缓存的新副本
+        """
+        new_cache = CustomQuantizeDynamicCache(self.cache_config)
+        if hasattr(self, '_quantized_key_cache'):
+            new_cache._quantized_key_cache = [k.clone() for k in self._quantized_key_cache]
+            new_cache._quantized_value_cache = [v.clone() for v in self._quantized_value_cache]
+        new_cache.key_cache = [k.clone() for k in self.key_cache]
+        new_cache.value_cache = [v.clone() for v in self.value_cache]
+        new_cache.cache_kwargs = copy.deepcopy(self.cache_kwargs)
+        new_cache.group_cache = copy.deepcopy(self.group_cache)
+        new_cache.meta = copy.deepcopy(self.meta)
+        new_cache._seen_tokens = self._seen_tokens
+        return new_cache
+
+
+def create_cache(quantize_nbits: Optional[int] = 0):
+    """
+    根据参数创建合适的缓存实例
+
+    Args:
+        args: 包含量化相关参数的命名空间对象
+
+    Returns:
+        CustomDynamicCache or CustomQuantizeDynamicCache: 缓存实例
+    """
+    if quantize_nbits > 0:
+        quan_cache_config = QuantizedCacheConfig(nbits=quantize_nbits)
+        return CustomQuantizeDynamicCache(quan_cache_config)
+    else:
+        return CustomDynamicCache()
+
+def manual_deepcopy_kv_cache(cache_obj):
+    """
+    Manually performs a deep copy of a custom KV cache object,
+    avoiding the issues with quanto's __deepcopy__.
+    """
+    # 1. 创建一个新的、空的 cache 对象实例
+    if isinstance(cache_obj, CustomQuantizeDynamicCache):
+        # 如果是量化缓存，需要传入配置
+        new_cache = CustomQuantizeDynamicCache(cache_obj.cache_config)
+    elif isinstance(cache_obj, CustomDynamicCache):
+        new_cache = CustomDynamicCache()
+    else:
+        # 如果有其他类型的缓存，可以在这里扩展
+        raise TypeError(f"Unsupported cache type for manual deepcopy: {type(cache_obj)}")
+
+    # 2. 复制非张量元数据
+    # 使用标准 deepcopy 是安全的，因为这些是字典和列表
+    new_cache.cache_kwargs = copy.deepcopy(cache_obj.cache_kwargs)
+    new_cache.meta = copy.deepcopy(cache_obj.meta)
+
+    # 3. 逐层复制核心的 key-value 张量缓存
+    if hasattr(cache_obj, '_quantized_key_cache'): # 处理 QuantoQuantizedCache
+        for layer_cache in cache_obj._quantized_key_cache:
+            # .clone().detach() 是安全复制张量的标准方法
+            new_cache._quantized_key_cache.append(layer_cache.clone().detach())
+        for layer_cache in cache_obj._quantized_value_cache:
+            new_cache._quantized_value_cache.append(layer_cache.clone().detach())
+        new_cache._seen_tokens = cache_obj._seen_tokens
+
+    if hasattr(cache_obj, 'key_cache'): # 处理 DynamicCache
+        for layer_cache in cache_obj.key_cache:
+            new_cache.key_cache.append(layer_cache.clone().detach())
+        for layer_cache in cache_obj.value_cache:
+            new_cache.value_cache.append(layer_cache.clone().detach())
+        new_cache._seen_tokens = cache_obj._seen_tokens
+
+    return new_cache
+
+def convert_tensor(data, cuda_device):
+    """转换结构体的 tensor device，如果cuda_device非 None，则将cpu 转换到 cuda，否则将 cuda 转换到 cpu
+    支持dict, list, tuple, set
+    """
+    
+    converted_count = [0]  # 使用列表以便在嵌套函数中修改
+    
+    def _convert_recursive(obj):
+        # 如果是torch tensor且在CUDA上
+        if torch.is_tensor(obj):
+            if cuda_device:
+                return obj.to(cuda_device) if obj.is_cpu else obj
+            if obj.is_cuda:
+                return obj.cpu()
+            return obj
+        
+        # 处理各种容器类型
+        elif isinstance(obj, dict):
+            return {k: _convert_recursive(v) for k, v in obj.items()}
+        
+        elif isinstance(obj, list):
+            return [_convert_recursive(item) for item in obj]
+        
+        elif isinstance(obj, tuple):
+            # 元组不可变，总是创建新的
+            return tuple(_convert_recursive(item) for item in obj)
+        
+        elif isinstance(obj, set):
+            return {_convert_recursive(item) for item in obj}
+        
+        # 其他数据类型直接返回
+        else:
+            return obj
+
+    return _convert_recursive(data)
+
+def copy_dict_to_cpu(d: dict):
+    ret = {}
+    for k, v in d.items():
+        ret[k] = v.cpu() if torch.is_tensor(v) and v.is_cuda else v
+    return ret
+
+def copy_dict_to_gpu(d: dict, device):
+    if not d:
+        return d
+    ret = {}
+    for k, v in d.items():
+        ret[k] = v.to(device) if torch.is_tensor(v) and not v.is_cuda else v
+    return ret
+
+def copy_kv_cache_to_device(cache_obj, cuda_device, copy_v: bool=True):
+    if isinstance(cache_obj, CustomQuantizeDynamicCache):
+        new_cache = CustomQuantizeDynamicCache(cache_obj.cache_config)
+    elif isinstance(cache_obj, CustomDynamicCache):
+        new_cache = CustomDynamicCache()
+    else:
+        raise TypeError(f"Unsupported cache type for manual deepcopy: {type(cache_obj)}")
+
+    # 复制非张量元数据
+    new_cache.cache_kwargs = convert_tensor(cache_obj.cache_kwargs, cuda_device)
+    new_cache.meta = convert_tensor(cache_obj.meta, cuda_device)
+
+    # 复制核心张量缓存
+    if hasattr(cache_obj, '_quantized_key_cache'):
+        new_cache._quantized_key_cache = convert_tensor(cache_obj._quantized_key_cache, cuda_device) 
+        if copy_v:
+            new_cache._quantized_value_cache = convert_tensor(cache_obj._quantized_value_cache, cuda_device) 
+        else:
+            new_cache._quantized_value_cache = cache_obj._quantized_value_cache
+        new_cache._seen_tokens = cache_obj._seen_tokens
+
+    if hasattr(cache_obj, 'key_cache'):
+        new_cache.key_cache = convert_tensor(cache_obj.key_cache, cuda_device)
+        if copy_v:
+            new_cache.value_cache = convert_tensor(cache_obj.value_cache, cuda_device)
+        else:
+            new_cache.value_cache = cache_obj.value_cache
+        new_cache._seen_tokens = cache_obj._seen_tokens
+
+    return new_cache
+
+# def copy_kv_cache_to_gpu(cache_obj, device, copy_v: bool=True):
+#     """
+#     手动执行缓存对象的深拷贝
+#     避免quanto库__deepcopy__的问题
+
+#     Args:
+#         cache_obj: 要拷贝的缓存对象
+
+#     Returns:
+#         缓存对象的深拷贝副本
+#     """
+#     if isinstance(cache_obj, CustomQuantizeDynamicCache):
+#         new_cache = CustomQuantizeDynamicCache(cache_obj.cache_config)
+#     elif isinstance(cache_obj, CustomDynamicCache):
+#         new_cache = CustomDynamicCache()
+#     else:
+#         raise TypeError(f"Unsupported cache type for manual deepcopy: {type(cache_obj)}")
+
+#     # 复制非张量元数据
+#     new_cache.cache_kwargs = copy_dict_to_gpu(cache_obj.cache_kwargs, device)
+#     new_cache.meta = copy_dict_to_gpu(cache_obj.meta, device)
+
+#     # 复制核心张量缓存
+#     if hasattr(cache_obj, '_quantized_key_cache'):
+#         new_cache._quantized_key_cache = [t.to(device) for t in cache_obj._quantized_key_cache]
+#         if copy_v:
+#             new_cache._quantized_value_cache = [t.to(device) for t in cache_obj._quantized_value_cache]
+#         else:
+#             new_cache._quantized_value_cache = cache_obj._quantized_value_cache
+#         new_cache._seen_tokens = cache_obj._seen_tokens
+
+#     if hasattr(cache_obj, 'key_cache'):
+#         new_cache.key_cache = [t.to(device) for t in cache_obj.key_cache]
+#         if copy_v:
+#             new_cache.value_cache = [t.to(device) for t in cache_obj.value_cache]
+#         else:
+#             new_cache.value_cache = cache_obj.value_cache
+
+#         new_cache._seen_tokens = cache_obj._seen_tokens
+
+#     return new_cache

src/utils/callbacks.py

@@ -0,0 +1,172 @@
+import os
+import time
+import typing
+# from deepspeed.accelerator import get_accelerator
+
+# NOTE: 最新版开始迁移到 integration_utils
+try:
+    from transformers.integrations import TrainerCallback
+except ImportError:
+    from transformers.integrations.integration_utils import TrainerCallback
+
+# USE_FLASH_ATTN, USE_XFORMERS_ATTN = False, False
+# if os.getenv('FLASH_ATTN', 'false').lower() == 'true':
+#     USE_FLASH_ATTN = True
+#     from mllm.utils.llama_flash_attn_monkey_patch import replace_llama_attn_with_flash_attn, restore_llama_attn
+# if os.getenv("XFORMERS_ATTN", 'false').lower() == 'true':
+#     USE_XFORMERS_ATTN = True
+#     from mllm.utils.llama_xformers_monkey_patch import replace_llama_attn_with_xformers_attn, restore_llama_attn
+
+class ModeltimeCallback(TrainerCallback):
+    def __init__(self):
+        self.model_time = 0.
+        self.data_time = 0.
+        self._start = 0.
+        self._end = 0.
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        self._end = time.time()
+    
+    def on_step_begin(self, args, state, control, **kwargs):
+        self._start = time.time()
+        self.data_time += (self._start - self._end)
+
+    def on_step_end(self, args, state, control, **kwargs):
+        self._end = time.time()
+        self.model_time += (self._end - self._start)
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if not state.is_world_process_zero:
+            return
+        
+        data_time = self.data_time / state.logging_steps
+        model_time = self.model_time / state.logging_steps
+        self.data_time = 0.
+        self.model_time = 0.
+        
+        info = f'\nSTEP: {state.global_step}, data_time: {data_time:.3f}, model_time: {model_time:.3f}'
+        print(info)
+
+
+class SacredCallback(TrainerCallback):
+    def __init__(self, _run=None):
+        self._run = _run
+        self.model_time = 0.
+        self.data_time = 0.
+        self._start = 0.
+        self._end = 0.
+        self._zero_loss_cnt = 0
+    
+    def on_train_begin(self, args, state, control, **kwargs):
+        if self._run:
+            self._end = time.time()
+    
+    def on_step_begin(self, args, state, control, **kwargs):
+        if self._run:
+            self._start = time.time()
+            self.data_time += (self._start - self._end)
+
+    def on_step_end(self, args, state, control, **kwargs):
+        if self._run:
+            self._end = time.time()
+            self.model_time += (self._end - self._start)
+    
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if not state.is_world_process_zero:
+            return
+        
+        if self._run is None:
+            return 
+
+        data_time = self.data_time / state.logging_steps
+        model_time = self.model_time / state.logging_steps
+        self.data_time = 0.
+        self.model_time = 0.
+        
+        self._run.log_scalar("loss", logs.get("loss", 0.), state.global_step)
+        self._run.log_scalar("learning_rate", logs.get("learning_rate", 0.), state.global_step)
+        self._run.log_scalar("grad_norm", logs.get("grad_norm", 0.), state.global_step)
+        self._run.log_scalar("epoch", state.epoch, state.global_step)
+        self._run.log_scalar("data_time", data_time, state.global_step)
+        self._run.log_scalar("model_time", model_time, state.global_step)
+        
+        if logs.get("loss", 0.) == 0.:
+            self._zero_loss_cnt += 1
+            if self._zero_loss_cnt > 1:
+                raise RuntimeError("Loss is zero, something is wrong!")
+        
+
+class ModelEvalCallback(TrainerCallback):
+    def __init__(self, _run=None, multitest=None, trainer=None, gen_kwargs=None):
+        self._run = _run
+        # datasets dict
+        # key1: dataset_name
+        # val1: {'dataset': dataset inst, 'compute_metric': metric inst}
+        self.multitest = typing.cast(dict, multitest)
+        self.trainer = trainer
+        self.gen_kwargs = gen_kwargs
+    
+    def on_step_end(self, args, state, control, **kwargs):
+        if args.eval_steps is None:
+            eval_steps = args.save_steps
+        elif isinstance(args.eval_steps, int) and args.eval_steps > 0:
+            eval_steps = args.eval_steps
+        else:
+            return
+        if state.global_step > 0 and state.global_step % eval_steps == 0:
+            if not args.do_multi_predict:
+                return
+            
+            # flash-attn currently not supports eval mode!
+            # if USE_FLASH_ATTN or USE_XFORMERS_ATTN:
+            #     restore_llama_attn()
+            
+            old_compute_metrics = self.trainer.compute_metrics
+            
+            for dataset_idx, (dataset_name, item) in enumerate(self.multitest.items()):
+                print(f'processing multitest set {dataset_idx}/{len(self.multitest)}: {dataset_name}')
+                _ds = item['dataset']
+                _compute_metrics = item['compute_metric']
+                _prefix = dataset_name
+                
+                self.trainer.compute_metrics = _compute_metrics
+                # transformers.trainer_utils.PredictionOutput
+                _pred_results = self.trainer.predict(_ds, metric_key_prefix=_prefix, **self.gen_kwargs)
+                if state.is_world_process_zero:
+                    self.trainer.log_metrics(_prefix, _pred_results.metrics)  # noqa
+                    self.trainer.save_metrics(_prefix, _pred_results.metrics)  # noqa
+                    self.trainer.save_prediction(_pred_results, file_key_prefix=_prefix)
+                    
+                    if self._run is not None:
+                        keywords_to_remove = ['runtime', 'second']
+                        for k, v in _pred_results.metrics.items():
+                            # remove time releated metrics
+                            if any(kw in k for kw in keywords_to_remove):
+                                continue
+                            self._run.log_scalar(f'{k}', v, state.global_step)
+            
+            self.trainer.compute_metrics = old_compute_metrics
+            
+            # if USE_FLASH_ATTN:
+            #     replace_llama_attn_with_flash_attn()
+            # if USE_XFORMERS_ATTN:
+            #     replace_llama_attn_with_xformers_attn()
+
+class DSEmptyCacheCallback(TrainerCallback):
+    def on_step_end(self, args, state, control, **kwargs):
+        empty_cache_steps = int(os.getenv("EMPTY_CACHE_STEP", '0').strip())
+        can_flush = state.global_step > 0 and empty_cache_steps > 0 and state.global_step % empty_cache_steps == 0
+        if can_flush:
+            # print('Flush Cache here.')
+            get_accelerator().empty_cache()
+            
+
+# usage: https://github.com/yqhu/profiler-workshop/blob/c8d4a7c30a61cc7b909d89f88f5fd36b70c55769/hf_training_trainer_prof.py#L49C6-L49C28
+# additionally, with_modules can be set True, with_flops must be set False
+# deps: pip install -U tensorboard-plugin-profilepip torch_tb_profiler 
+class ProfCallback(TrainerCallback):
+    def __init__(self, prof):
+        self.prof = prof
+
+    def on_step_end(self, args, state, control, **kwargs):
+        self.prof.step()

src/utils/common.py

@@ -0,0 +1,187 @@
+import copy
+from typing import List, Union, Dict
+
+import PIL.Image
+import torch
+import numpy as np
+import torchvision.transforms.functional as F
+import transformers
+
+from transformers import PreTrainedTokenizer
+
+IGNORE_INDEX = -100
+
+def print_trainable_params(model: torch.nn.Module) -> None:
+    trainable_params, all_param = 0, 0
+    for param in model.parameters():
+        num_params = param.numel()
+        # if using DS Zero 3 and the weights are initialized empty
+        if num_params == 0 and hasattr(param, "ds_numel"):
+            num_params = param.ds_numel
+        all_param += num_params
+        if param.requires_grad:
+            trainable_params += num_params
+    print("trainable params: {:d} || all params: {:d} || trainable%: {:.4f}".format(
+        trainable_params, all_param, 100 * trainable_params / all_param))
+
+
+def post_process_generate_ids(tokenizer: PreTrainedTokenizer, ids: torch.Tensor):
+    ids = copy.deepcopy(ids)  # do not modify origin preds and targets
+    ids[ids < 0] = tokenizer.pad_token_id
+    # pad_to_multiof 开启后, 多余的部分没法解码, 这里暂时替换为 ','
+    ids[ids >= len(tokenizer)] = tokenizer.convert_tokens_to_ids(',')
+    return ids
+
+
+def decode_generate_ids(tokenizer: PreTrainedTokenizer, ids: torch.Tensor) -> Union[List[str], str]:
+    assert ids.ndim in [1, 2]
+    only_one_sentence = ids.ndim == 1
+    if only_one_sentence:
+        ids = ids.unsqueeze(0)
+    ids = post_process_generate_ids(tokenizer, ids)
+    res = tokenizer.batch_decode(ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+    if only_one_sentence:
+        return res[0]
+    return res
+
+
+
+def draw_bounding_boxes(
+        image: Union[torch.Tensor, PIL.Image.Image],
+        boxes: Union[torch.Tensor, List, np.ndarray],
+        **kwargs,
+):
+    if isinstance(image, PIL.Image.Image):
+        from torchvision.transforms import PILToTensor
+        image = PILToTensor()(image)
+    assert isinstance(image, torch.Tensor), ""
+
+    if not isinstance(boxes, torch.Tensor):
+        boxes = torch.as_tensor(boxes)
+    assert isinstance(boxes, torch.Tensor)
+
+    from torchvision.utils import draw_bounding_boxes as _draw_bounding_boxes
+    return _draw_bounding_boxes(image, boxes, **kwargs)
+
+
+# https://github.com/huggingface/tokenizers/issues/247#issuecomment-675458087
+def smart_tokenizer_and_embedding_resize(
+        special_tokens_dict: Dict,
+        tokenizer: transformers.PreTrainedTokenizer,
+        model: transformers.PreTrainedModel,
+):
+    """Resize tokenizer and embedding.
+
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+def patch_transformer_logging():
+    import logging
+    import transformers
+    def enable_explicit_format():
+        handlers = transformers.utils.logging._get_library_root_logger().handlers
+
+        for handler in handlers:
+            formatter = logging.Formatter("[(%(levelname)s) %(pathname)s:%(lineno)s ] %(asctime)s >> %(message)s")
+            handler.setFormatter(formatter)
+    transformers.utils.logging.enable_explicit_format = enable_explicit_format
+
+def print_model_stats(model):
+    """
+    通用模型参数统计工具。
+    自动识别是 Dense 还是 MoE 模型，并计算 Total vs Active 参数量。
+    """
+    
+    # 1. 计算物理总参数量 (Total Parameters)
+    # 使用 set 避免计算共享参数 (Shared Weights)，例如 Embedding 和 lm_head 共享权重的情况
+    unique_params = {p.data_ptr(): p for p in model.parameters()}.values()
+    total_params = sum(p.numel() for p in unique_params)
+    
+    # 2. 初始化激活参数量 (Active Parameters)
+    # 默认假设是 Dense 模型，所有参数都是激活的
+    active_params = total_params
+    
+    moe_infos = [] # 用于存储发现的 MoE 层信息
+    
+    # 3. 遍历所有子模块，寻找 MoE 层特征
+    # 我们不匹配类名，而是匹配"特征" (Duck Typing)
+    for name, module in model.named_modules():
+        # 特征判定：有 num_experts 属性，且有一个叫 experts 的 ModuleList
+        if hasattr(module, 'num_experts') and hasattr(module, 'experts') and isinstance(module.experts, nn.ModuleList):
+            
+            # 获取关键超参
+            num_experts = getattr(module, 'num_experts', 0)
+            # 兼容不同的 top_k 命名 (top_k 或 num_experts_per_tok)
+            top_k = getattr(module, 'top_k', getattr(module, 'num_experts_per_tok', 0))
+            
+            # 如果找不到 top_k，可能不是标准的 Sparse MoE，跳过
+            if top_k == 0: 
+                continue
+
+            # --- 核心计算逻辑 ---
+            # 1. 计算单个专家的参数量 (假设所有专家结构相同，取第一个)
+            # 这里必须用 recursion=True 确保统计专家内部所有层
+            single_expert_params = sum(p.numel() for p in module.experts[0].parameters())
+            
+            # 2. 计算"休眠"专家数量
+            dormant_experts = num_experts - top_k
+            
+            # 3. 从激活总数中扣除休眠专家的参数
+            # 注意：total_params 里已经包含了 N 个专家，我们只需要减去 (N-K) 个
+            if dormant_experts > 0:
+                deduction = dormant_experts * single_expert_params
+                active_params -= deduction
+                
+                moe_infos.append({
+                    "layer": name,
+                    "experts": num_experts,
+                    "active": top_k,
+                    "expert_size": single_expert_params
+                })
+
+    # --- 4. 格式化输出 ---
+    def format_num(num):
+        if num >= 1e9: return f"{num/1e9:.2f}B"
+        if num >= 1e6: return f"{num/1e6:.2f}M"
+        if num >= 1e3: return f"{num/1e3:.2f}K"
+        return str(num)
+
+    print("=" * 50)
+    print(f"Model Architecture Analysis")
+    print("=" * 50)
+    
+    if len(moe_infos) > 0:
+        print(f"👉 Detection: MoE Model (Sparse Mixture-of-Experts)")
+        print(f"   - Found {len(moe_infos)} MoE layers")
+        print(f"   - Config: {moe_infos[0]['experts']} Experts, Top-{moe_infos[0]['active']} Active")
+    else:
+        print(f"👉 Detection: Dense Model (Standard Transformer)")
+        
+    print("-" * 50)
+    print(f"Total Parameters (VRAM):  {format_num(total_params)}")
+    print(f"Active Parameters (FLOPs): {format_num(active_params)}")
+    
+    if len(moe_infos) > 0:
+        sparsity = 1 - (active_params / total_params)
+        print(f"Sparsity Ratio:            {sparsity:.2%}")
+        # 计算相比 Dense 版本的倍数
+        # 假设 Dense 版本就是 active_params 大小（不太严谨但直观）
+        print(f"Upcycling Scale:           {total_params/active_params:.2f}x Larger than Dense Base")
+    else:
+        print(f"Sparsity Ratio:            0.00% (Dense)")
+        
+    print("=" * 50)
+    
+    return total_params, active_params

src/utils/data_utils.py

@@ -0,0 +1,318 @@
+import os
+import lmdb
+from multiprocessing import Pool
+from tqdm import tqdm
+import importlib.metadata
+import importlib.util
+from packaging import version
+from typing import TYPE_CHECKING
+from functools import lru_cache
+from datasets import Dataset, IterableDataset
+from src.utils.common import IGNORE_INDEX
+from collections import defaultdict
+from functools import partial
+import bisect
+from typing import List, Sequence, Tuple, Optional, Union
+from src.utils.common import pdb_debug
+
+if TYPE_CHECKING:
+    from packaging.version import Version
+
+def write_lmdb(output_dir, name):
+
+    os.makedirs(output_dir, exist_ok=True)
+    output_name = os.path.join(output_dir, f'{name}.lmdb')
+
+    try:
+        os.remove(output_name)
+    except:
+        pass
+    env_new = lmdb.open(
+        output_name,
+        subdir=False,
+        readonly=False,
+        lock=False,
+        readahead=False,
+        meminit=False,
+        max_readers=1,
+        map_size=int(100e9),
+    )
+    txn_write = env_new.begin(write=True)
+
+    return txn_write, env_new
+
+def read_lmdb(lmdb_path):
+    env = lmdb.open(
+        lmdb_path,
+        subdir=False,
+        readonly=True,
+        lock=False,
+        readahead=False,
+        meminit=False,
+        max_readers=256,
+    )
+    txn = env.begin()
+    return env, txn
+
+
+def get_length(index):
+    return index, len(global_dataset[index]["input_ids"])
+
+def get_sequence_length(dataset, num_worker=16):
+    global global_dataset
+    global_dataset = dataset
+    num_data = len(dataset)
+    lengths = [0] * num_data
+    with Pool(processes=num_worker) as pool:
+        iters = pool.imap(get_length, range(num_data))
+        for i, length in tqdm(iters, total=num_data):
+            lengths[i] = length
+    return lengths
+
+def get_data(index):
+    item = global_torch_dataset[index]
+    length = len(global_torch_dataset[index]["input_ids"])
+    return item, index, length
+
+def torch_dataset_to_hf_dataset(torch_dataset, num_worker=16):
+    global global_torch_dataset
+    global_torch_dataset = torch_dataset
+    num_data = len(global_torch_dataset)
+    lengths = [0] * num_data
+    hf_dict = {key: [] for key in torch_dataset[0].keys()}
+    with Pool(processes=num_worker) as pool:
+        iters = pool.imap(get_data, range(num_data))
+        for data, i, length in tqdm(iters, total=num_data):
+            for key, value in data.items():
+                hf_dict[key].append(value)
+            lengths[i] = length
+    hf_dataset = Dataset.from_dict(hf_dict)
+    return hf_dataset, lengths
+
+def _get_package_version(name: str) -> "Version":
+    try:
+        return version.parse(importlib.metadata.version(name))
+    except Exception:
+        return version.parse("0.0.0")
+
+@lru_cache
+def is_transformers_version_greater_than(content: str):
+    return _get_package_version("transformers") >= version.parse(content)
+
+@lru_cache
+def is_transformers_version_equal_to_4_46():
+    return version.parse("4.46.0") <= _get_package_version("transformers") <= version.parse("4.46.1")
+
+def search_for_fit(numbers: Sequence[int], capacity: int) -> int:
+    r"""
+    Finds the index of largest number that fits into the knapsack with the given capacity.
+    """
+    index = bisect.bisect(numbers, capacity)
+    return -1 if index == 0 else (index - 1)
+
+def greedy_knapsack(numbers: List[int], capacity: int) -> List[List[int]]:
+    r"""
+    An efficient greedy algorithm with binary search for the knapsack problem.
+    """
+    numbers.sort()  # sort numbers in ascending order for binary search
+    knapsacks = []
+
+    while numbers:
+        current_knapsack = []
+        remaining_capacity = capacity
+
+        while True:
+            index = search_for_fit(numbers, remaining_capacity)
+            if index == -1:
+                break  # no more numbers fit in this knapsack
+
+            remaining_capacity -= numbers[index]  # update the remaining capacity
+            current_knapsack.append(numbers.pop(index))  # add the number to knapsack
+
+        knapsacks.append(current_knapsack)
+
+    return knapsacks
+
+def preprocess_packed_supervised_dataset(examples, tokenizer, cutoff_len):
+    valid_num = 0
+    batch_input_ids, batch_labels = [], []
+    lengths = []
+    length2indexes = defaultdict(list)
+    for i in range(len(examples["input_ids"])):
+        input_ids, labels = examples["input_ids"][i], examples["labels"][i]
+        length = len(input_ids)
+        if length >= cutoff_len - 1:
+            continue
+        else:
+            lengths.append(length)
+            length2indexes[length].append(valid_num)
+            batch_input_ids.append(input_ids)
+            batch_labels.append(labels)
+            valid_num += 1
+    model_inputs = defaultdict(list)
+    knapsacks = greedy_knapsack(lengths, cutoff_len - 1)
+    for knapsack in knapsacks:
+        packed_input_ids, packed_attention_masks, packed_labels = [], [], []
+        for i, length in enumerate(knapsack):
+            index = length2indexes[length].pop()
+            packed_input_ids += batch_input_ids[index]
+            packed_labels += batch_labels[index]
+            packed_attention_masks += [1] * len(batch_input_ids[index])
+
+        if len(packed_input_ids) < cutoff_len:
+            pad_length = cutoff_len - len(packed_input_ids)
+            packed_input_ids += [tokenizer.pad_token_id] * pad_length
+            packed_labels += [IGNORE_INDEX] * pad_length
+            packed_attention_masks += [1] * pad_length  # more efficient flash_attn
+
+        if len(packed_input_ids) != cutoff_len:
+            raise ValueError("The length of packed example should be identical to the cutoff length.")
+
+        model_inputs["input_ids"].append(packed_input_ids)
+        model_inputs["attention_mask"].append(packed_attention_masks)
+        model_inputs["position_ids"].append(list(range(len(packed_input_ids))))
+        model_inputs["labels"].append(packed_labels)
+    return model_inputs
+
+def pad_sequence(examples, cutoff_len, tokenizer):
+    max_length = cutoff_len
+    input_pad_token_id = tokenizer.pad_token_id
+    label_pad_token_id = IGNORE_INDEX
+
+    for k, v in examples.items():
+        if k.endswith("input_ids"):
+            pad_token_id = input_pad_token_id
+        elif k.endswith("labels"):
+            pad_token_id = label_pad_token_id
+            # shift labels here
+            for i in range(len(v)):
+                v[i] = v[i][1:]
+        elif k.endswith("attention_mask"):
+            pad_token_id = 0
+        elif k.endswith("position_ids"):
+            pad_token_id = max_length - 1  # pad the max position id
+        elif k == "images" or k == "videos":
+            pad_token_id = -1
+            continue  # TODO: haven't tested multi-modal yet
+        else:
+            continue
+        for i in range(len(v)):
+            v[i].extend([pad_token_id] * (max_length - len(v[i])))
+        examples[k] = v
+
+    return examples
+
+def preprocess_sp_dataset(seq_ids, world_size, sequence_parallel_mode):
+    if sequence_parallel_mode == "zigzag-ring":
+        step = len(seq_ids) // (2 * world_size)
+        value_chunks = [seq_ids[s : s + step] for s in range(0, len(seq_ids), step)]
+        local_values = list()
+        for rank in range(world_size):
+            local_values.append(value_chunks[rank] + value_chunks[2 * world_size - rank - 1])
+        return local_values
+    elif sequence_parallel_mode == "ulysses":
+        step = len(seq_ids) // world_size
+        local_values = [seq_ids[s : s + step] for s in range(0, len(seq_ids), step)]
+        return local_values
+    else:
+        raise NotImplementedError("Other sequence parallel modes are to be implemented.")
+
+
+# sp for Sequence Parallel
+def sp_split(examples, sequence_parallel_size, sequence_parallel_mode="ulysses"):
+    for k, v in examples.items():
+        chunks = list()
+        for row in v:
+            if k.endswith("attention_mask"):
+                chunks.extend([row] * sequence_parallel_size)
+            elif row is None:
+                chunks.extend([None] * sequence_parallel_size)
+            else:
+                chunks.extend(
+                    preprocess_sp_dataset(row, sequence_parallel_size, sequence_parallel_mode)
+                )
+        examples[k] = chunks
+    return examples
+
+def get_sequence_parallel_preprocess(stage, tokenizer, cutoff_len=None, sequence_parallel_size=1, sequence_parallel_mode="ulysses"):
+    if stage == "pad":
+        assert cutoff_len is not None
+        preprocess_func = partial(pad_sequence, cutoff_len=cutoff_len, tokenizer=tokenizer)
+    elif stage == "split":
+        preprocess_func = partial(sp_split, sequence_parallel_size=sequence_parallel_size, sequence_parallel_mode=sequence_parallel_mode)
+    else:
+        raise NotImplementedError(f"Unexpected stage in sequence_parallel_preprocess: {stage}")
+
+    return preprocess_func
+
+def _get_sequence_parallel_dataset(dataset, num_works, tokenizer=None, cutoff_len=10000, 
+                                   sequence_parallel_size=1, sequence_parallel_mode="ulysses", 
+                                   cache_dataset_overwrite=False) -> Optional[Union["Dataset", "IterableDataset"]]:
+    kwargs = dict(
+        num_proc=num_works,
+        load_from_cache_file=not cache_dataset_overwrite,
+        desc="Running padding split on dataset",
+    )
+    pad_sequence_func = get_sequence_parallel_preprocess(
+        stage="pad", 
+        tokenizer=tokenizer, 
+        cutoff_len=cutoff_len
+    )
+    padded_dataset = dataset.map(
+        pad_sequence_func, batched=True, batch_size=num_works, **kwargs
+    )
+    kwargs = dict(
+        num_proc=num_works,
+        load_from_cache_file=not cache_dataset_overwrite,
+        desc="Running sequence parallel split on dataset",
+    )
+    sp_dataset_func = get_sequence_parallel_preprocess(
+        stage="split", 
+        tokenizer=tokenizer, 
+        sequence_parallel_size=sequence_parallel_size, 
+        sequence_parallel_mode=sequence_parallel_mode,
+    )
+    sp_dataset = padded_dataset.map(
+        sp_dataset_func, batched=True, batch_size=num_works, **kwargs
+    )
+    return sp_dataset
+
+def packing_dataset(dataset, tokenizer, cutoff_len, num_worker, cache_dataset_overwrite):
+    preprocess_func = partial(preprocess_packed_supervised_dataset, tokenizer=tokenizer, cutoff_len=cutoff_len)
+    kwargs = dict(
+        num_proc=num_worker,
+        load_from_cache_file=not cache_dataset_overwrite,
+        desc="Running postprocess on dataset",
+    )
+    import pdb; pdb.set_trace()
+    dataset = dataset.map(
+        preprocess_func,
+        batched=True,
+        batch_size=num_worker,
+        **kwargs,
+    )
+    return dataset
+
+def data_post_process_sequence_parallel(
+    dataset, 
+    training_args, 
+    sequence_parallel_size, 
+    sequence_parallel_mode, 
+    cutoff_len, 
+    num_worker=16, 
+    packing=False, 
+    tokenizer=None,
+    cache_dataset_overwrite=False,
+):
+    dataset = dataset.shuffle(seed=training_args.seed)
+    if packing:
+        dataset = packing_dataset(dataset, tokenizer, cutoff_len, num_worker, cache_dataset_overwrite)
+
+    dataset = _get_sequence_parallel_dataset(dataset, 
+                                             num_works=num_worker, 
+                                             tokenizer=tokenizer, 
+                                             cutoff_len=cutoff_len, 
+                                             sequence_parallel_size=sequence_parallel_size,
+                                             sequence_parallel_mode=sequence_parallel_mode,
+                                             cache_dataset_overwrite=cache_dataset_overwrite)
+    return dataset

src/utils/gpu_monitor.py

@@ -0,0 +1,384 @@
+import pynvml
+import time
+import threading
+from datetime import datetime
+import sys
+import os
+
+class GPUMemoryMonitor:
+    def __init__(self, gpu_index=0, interval=1.0, unit='GB'):
+        """
+        初始化GPU显存监控器
+        
+        Args:
+            gpu_index: 要监控的GPU索引，默认0
+            interval: 监控间隔时间（秒），默认1秒
+            unit: 返回的单位，支持 'MB' 或 'GB'，默认'GB'
+        """
+        self.gpu_index = gpu_index
+        self.interval = interval
+        self.unit = unit.upper()
+        
+        # 验证单位参数
+        if self.unit not in ['MB', 'GB']:
+            raise ValueError("单位必须是 'MB' 或 'GB'")
+        
+        # 监控相关状态
+        self.monitor_thread = None
+        self._running = False
+        self._lock = threading.Lock()
+        self.peak_memory_usage = 0  # 峰值显存使用量
+        self.start_time = None
+        self.stop_time = None
+        
+        # 初始化NVML
+        try:
+            pynvml.nvmlInit()
+            self.device_count = pynvml.nvmlDeviceGetCount()
+            if self.gpu_index >= self.device_count:
+                raise ValueError(f"GPU索引 {self.gpu_index} 超出范围，系统只有 {self.device_count} 个GPU")
+            self.handle = pynvml.nvmlDeviceGetHandleByIndex(self.gpu_index)
+            
+            # 获取GPU名称
+            self.gpu_name = pynvml.nvmlDeviceGetName(self.handle)
+            
+        except Exception as e:
+            print(f"初始化NVML失败: {e}")
+            if 'pynvml' in sys.modules:
+                pynvml.nvmlShutdown()
+            raise
+    
+    def _get_memory_info(self):
+        """获取当前GPU显存信息"""
+        try:
+            mem_info = pynvml.nvmlDeviceGetMemoryInfo(self.handle)
+            return mem_info
+        except Exception as e:
+            print(f"获取GPU {self.gpu_index} 显存信息失败: {e}")
+            return None
+    
+    def _convert_units(self, bytes_value):
+        """转换字节为指定单位"""
+        if self.unit == 'MB':
+            return bytes_value / (1024 * 1024)
+        else:  # GB
+            return bytes_value / (1024 * 1024 * 1024)
+    
+    def _get_current_usage(self):
+        """获取当前显存使用量（按指定单位）"""
+        mem_info = self._get_memory_info()
+        if mem_info:
+            return self._convert_units(mem_info.used)
+        return 0
+    
+    def _monitor_loop(self):
+        """监控循环，运行在单独线程中"""
+        # print(f"开始监控 GPU {self.gpu_index} ({self.gpu_name})")
+        # print(f"监控间隔: {self.interval}秒")
+        # print(f"单位: {self.unit}")
+        # print("按回车键停止监控...\n")
+        
+        try:
+            while self._running:
+                # 获取当前显存使用量
+                current_usage = self._get_current_usage()
+                
+                # 更新峰值
+                with self._lock:
+                    if current_usage > self.peak_memory_usage:
+                        self.peak_memory_usage = current_usage
+                
+                # 等待下一个监控周期
+                time.sleep(self.interval)
+                
+        except Exception as e:
+            print(f"监控线程出错: {e}")
+        finally:
+            # print(f"GPU {self.gpu_index} 监控线程结束")
+            pass
+    
+    def start(self):
+        """
+        启动显存监控
+        
+        Returns:
+            bool: 是否成功启动
+        """
+        if self._running:
+            # print(f"GPU {self.gpu_index} 监控已经在运行中")
+            return False
+        
+        try:
+            # 重置峰值数据
+            with self._lock:
+                self.peak_memory_usage = 0
+            
+            # 记录开始时间
+            self.start_time = datetime.now()
+            
+            # 启动监控线程
+            self._running = True
+            self.monitor_thread = threading.Thread(target=self._monitor_loop)
+            self.monitor_thread.daemon = True
+            self.monitor_thread.start()
+            
+            # 验证线程已启动
+            time.sleep(0.1)
+            if self.monitor_thread.is_alive():
+                # print(f"GPU {self.gpu_index} 监控已启动")
+                return True
+            else:
+                self._running = False
+                # print(f"GPU {self.gpu_index} 监控线程启动失败")
+                return False
+                
+        except Exception as e:
+            print(f"启动GPU {self.gpu_index} 监控失败: {e}")
+            self._running = False
+            return False
+    
+    def stop(self, verbose=False):
+        """
+        停止显存监控并返回峰值显存使用量
+        
+        Args:
+            verbose: 是否打印详细信息
+            
+        Returns:
+            float: 峰值显存使用量（单位：GB或MB）
+        """
+        if not self._running:
+            if verbose:
+                print(f"GPU {self.gpu_index} 监控未在运行")
+            return 0.0
+        
+        try:
+            # 停止监控线程
+            self._running = False
+            self.stop_time = datetime.now()
+            
+            # 等待线程结束（最多等待2秒）
+            if self.monitor_thread and self.monitor_thread.is_alive():
+                self.monitor_thread.join(timeout=2.0)
+            
+            # 获取最终峰值
+            peak_usage = 0.0
+            with self._lock:
+                peak_usage = self.peak_memory_usage
+            
+            if verbose:
+                duration = (self.stop_time - self.start_time).total_seconds()
+                print(f"\n{'='*50}")
+                print(f"GPU {self.gpu_index} ({self.gpu_name}) 监控结果")
+                print(f"{'='*50}")
+                print(f"开始时间: {self.start_time.strftime('%Y-%m-%d %H:%M:%S')}")
+                print(f"结束时间: {self.stop_time.strftime('%Y-%m-%d %H:%M:%S')}")
+                print(f"监控时长: {duration:.1f} 秒")
+                print(f"峰值显存使用量: {peak_usage:.3f} {self.unit}")
+                
+                # 获取当前显存信息作为对比
+                mem_info = self._get_memory_info()
+                if mem_info:
+                    total = self._convert_units(mem_info.total)
+                    current = self._convert_units(mem_info.used)
+                    print(f"当前显存使用: {current:.3f} / {total:.3f} {self.unit}")
+                    print(f"峰值占比: {(peak_usage/total*100):.1f}%")
+                print(f"{'='*50}")
+            
+            return peak_usage
+            
+        except Exception as e:
+            print(f"停止GPU {self.gpu_index} 监控时出错: {e}")
+            return 0.0
+    
+    def get_current_usage(self):
+        """获取当前显存使用量（不停止监控）"""
+        return self._get_current_usage()
+    
+    def get_peak_usage(self):
+        """获取当前记录的峰值显存使用量（不停止监控）"""
+        with self._lock:
+            return self.peak_memory_usage
+    
+    def is_running(self):
+        """检查监控是否在运行"""
+        return self._running
+    
+    def __del__(self):
+        """析构函数，确保清理资源"""
+        if self._running:
+            self.stop(verbose=False)
+        try:
+            pynvml.nvmlShutdown()
+        except:
+            pass
+
+
+# 使用示例函数
+def monitor_gpu_memory_example():
+    """使用示例"""
+    print("GPU显存监控示例")
+    print("=" * 50)
+    
+    try:
+        # 创建监控器
+        monitor = GPUMemoryMonitor(
+            gpu_index=0,      # 监控第一个GPU
+            interval=0.5,     # 每0.5秒检查一次
+            unit='GB'         # 使用GB作为单位
+        )
+        
+        # 启动监控
+        if monitor.start():
+            # 这里可以运行你的GPU任务
+            print("\n现在开始运行你的GPU任务...")
+            print("监控正在后台进行")
+            
+            # 模拟一些工作（在实际使用中，这里应该是你的GPU任务）
+            print("按回车键停止监控并获取峰值显存用量...")
+            input()  # 等待用户按回车
+            
+            # 停止监控并获取结果
+            peak_memory_gb = monitor.stop()
+            
+            print(f"\n监控完成！峰值显存用量: {peak_memory_gb:.3f} GB")
+    
+    except Exception as e:
+        print(f"错误: {e}")
+
+
+# 多GPU监控示例
+class MultiGPUMonitor:
+    """多GPU监控器"""
+    def __init__(self, gpu_indices=None, interval=1.0, unit='GB'):
+        """
+        初始化多GPU监控器
+        
+        Args:
+            gpu_indices: 要监控的GPU索引列表，None表示监控所有GPU
+            interval: 监控间隔
+            unit: 返回的单位
+        """
+        pynvml.nvmlInit()
+        device_count = pynvml.nvmlDeviceGetCount()
+        
+        if gpu_indices is None:
+            gpu_indices = list(range(device_count))
+        
+        self.monitors = []
+        for idx in gpu_indices:
+            if idx < device_count:
+                monitor = GPUMemoryMonitor(gpu_index=idx, interval=interval, unit=unit)
+                self.monitors.append(monitor)
+            else:
+                print(f"警告: GPU索引 {idx} 不存在，跳过")
+    
+    def start_all(self):
+        """启动所有GPU监控"""
+        results = []
+        for monitor in self.monitors:
+            success = monitor.start()
+            results.append((monitor.gpu_index, success))
+        return results
+    
+    def stop_all(self):
+        """停止所有GPU监控并返回结果"""
+        results = {}
+        max_peak = 0.0
+        for monitor in self.monitors:
+            peak = monitor.stop()
+            results[monitor.gpu_index] = {
+                'name': monitor.gpu_name,
+                'peak_memory': peak,
+                'unit': monitor.unit
+            }
+            if peak > max_peak:
+                max_peak = peak
+        return results, max_peak
+    
+    def stop_all_and_get_max(self):
+        """停止所有监控并返回最大峰值"""
+        results = self.stop_all()
+        if not results:
+            return 0.0
+        
+        max_peak = max(item['peak_memory'] for item in results.values())
+        return max_peak
+
+
+# 快速使用函数
+def quick_monitor(gpu_index=0, interval=0.5, unit='GB', wait_for_input=True):
+    """
+    快速启动监控的便捷函数
+    
+    Args:
+        gpu_index: GPU索引
+        interval: 监控间隔
+        unit: 单位
+        wait_for_input: 是否等待用户输入
+    
+    Returns:
+        float: 峰值显存使用量
+    """
+    monitor = GPUMemoryMonitor(gpu_index=gpu_index, interval=interval, unit=unit)
+    
+    try:
+        if monitor.start():
+            if wait_for_input:
+                print("按回车键停止监控...")
+                input()
+            else:
+                # 如果不等待用户输入，这里可以设置其他停止条件
+                # 例如：监控特定时间或直到某个条件满足
+                print("监控已启动，将在后台运行")
+                print("调用 monitor.stop() 来停止并获取结果")
+                return monitor
+        return None
+    except Exception as e:
+        print(f"监控失败: {e}")
+        return None
+
+
+# 测试代码
+if __name__ == "__main__":
+    # # 示例1: 基本使用
+    # print("示例1: 基本使用")
+    # monitor = GPUMemoryMonitor(gpu_index=0, interval=0.5, unit='GB')
+    # monitor.start()
+    
+    # # 模拟一些GPU工作
+    # print("模拟GPU工作...")
+    # time.sleep(3)
+    
+    # # 停止并获取结果
+    # peak = monitor.stop()
+    # print(f"峰值显存: {peak:.3f} GB\n")
+    
+    # # 示例2: 使用便捷函数
+    # print("示例2: 使用便捷函数")
+    # result = quick_monitor(gpu_index=0, interval=0.2, unit='MB')
+    # if result:
+    #     # 这里 result 是 monitor 对象
+    #     time.sleep(2)
+    #     peak = result.stop()
+    #     print(f"峰值显存: {peak:.3f} MB")
+    
+    # 示例3: 多GPU监控
+    print("\n示例3: 多GPU监控")
+    try:
+        pynvml.nvmlInit()
+        device_count = pynvml.nvmlDeviceGetCount()
+        print(f"系统中有 {device_count} 个GPU")
+        
+        if device_count > 1:
+            # multi_monitor = MultiGPUMonitor(gpu_indices=[0, 1], interval=0.5, unit='GB')
+            multi_monitor = MultiGPUMonitor(interval=5, unit='GB')
+            multi_monitor.start_all()
+            time.sleep(2)
+            results, max_peak = multi_monitor.stop_all()
+            
+            for gpu_idx, data in results.items():
+                print(f"GPU {gpu_idx} ({data['name']}): {data['peak_memory']:.3f} {data['unit']}")
+            print(f"最大峰值显存使用量: {max_peak:.3f} GB")
+    except:
+        pass

src/utils/gpu_worker.py

@@ -0,0 +1,19 @@
+import os
+import torch
+
+class GpuWorker:
+
+    def __init__(self, gpu_id: int, envs: dict):
+        self.gpu_id = gpu_id
+        self._setup_environment(envs)
+        self.device = self._setup_device()
+
+    def _setup_environment(self, envs):
+        """设置环境变量"""
+        for k, v in envs.items():
+            os.environ[k] = v
+
+    def _setup_device(self):
+        """设置CUDA设备"""
+        torch.cuda.set_device(self.gpu_id)
+        return f"cuda:{self.gpu_id}"

src/utils/misc.py

@@ -0,0 +1,175 @@
+import os
+import subprocess
+import math
+import torch
+import torch.distributed as dist
+import transformers
+from typing import Union, Iterable, List, Dict, Tuple, Optional
+
+class DictWithDotAccess(dict):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        for arg in args:
+            if isinstance(arg, dict):
+                for key, value in arg.items():
+                    if isinstance(value, dict):
+                        value = DictWithDotAccess(value)
+                    self[key] = value
+
+        if kwargs:
+            for key, value in kwargs.items():
+                if isinstance(value, dict):
+                    value = DictWithDotAccess(value)
+                self[key] = value
+
+    def __getattr__(self, attr):
+        return self.get(attr)
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+def get_local_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return int(os.environ["LOCAL_RANK"])
+
+def Print(*args):
+    if get_rank() == 0:
+        print(*args)
+
+def get_sha():
+    cwd = os.path.dirname(os.path.abspath(__file__))
+
+    def _run(command):
+        return subprocess.check_output(command, cwd=cwd).decode("ascii").strip()
+
+    sha = "N/A"
+    diff = "clean"
+    branch = "N/A"
+    try:
+        sha = _run(["git", "rev-parse", "HEAD"])
+        subprocess.check_output(["git", "diff"], cwd=cwd)
+        diff = _run(["git", "diff-index", "HEAD"])
+        diff = "has uncommited changes" if diff else "clean"
+        branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])
+    except Exception:
+        pass
+    message = f"sha: {sha}, status: {diff}, branch: {branch}"
+    return message
+
+def patch_cosine_with_warmup_schedule(minimal_lr=0.0):
+    def _get_cosine_schedule_with_warmup_lr_lambda(
+        current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float
+    ):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+        return max(minimal_lr, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))
+    
+    transformers.optimization._get_cosine_schedule_with_warmup_lr_lambda = _get_cosine_schedule_with_warmup_lr_lambda
+    
+
+# patch torch clip grad norm so that we can skip nan grad norm
+from torch import Tensor, inf
+from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype, _has_foreach_support
+_tensor_or_tensors = Union[torch.Tensor, Iterable[torch.Tensor]]
+def clip_grad_norm_(
+        parameters: _tensor_or_tensors, max_norm: float, norm_type: float = 2.0,
+        error_if_nonfinite: bool = False, foreach: Optional[bool] = None) -> torch.Tensor:
+    r"""Clips gradient norm of an iterable of parameters.
+
+    The norm is computed over all gradients together, as if they were
+    concatenated into a single vector. Gradients are modified in-place.
+
+    Args:
+        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+            single Tensor that will have gradients normalized
+        max_norm (float): max norm of the gradients
+        norm_type (float): type of the used p-norm. Can be ``'inf'`` for
+            infinity norm.
+        error_if_nonfinite (bool): if True, an error is thrown if the total
+            norm of the gradients from :attr:`parameters` is ``nan``,
+            ``inf``, or ``-inf``. Default: False (will switch to True in the future)
+        foreach (bool): use the faster foreach-based implementation.
+            If ``None``, use the foreach implementation for CUDA and CPU native tensors and silently
+            fall back to the slow implementation for other device types.
+            Default: ``None``
+
+    Returns:
+        Total norm of the parameter gradients (viewed as a single vector).
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    grads = [p.grad for p in parameters if p.grad is not None]
+    max_norm = float(max_norm)
+    norm_type = float(norm_type)
+    if len(grads) == 0:
+        return torch.tensor(0.)
+    
+    if torch.isnan(max_norm) or torch.isinf(max_norm):
+        for grad in grads:
+            grad.zero_()
+        print('>>>Found nan or inf max_norm, set grads to zero')
+        return torch.tensor(0.)
+    
+    first_device = grads[0].device
+    grouped_grads: Dict[Tuple[torch.device, torch.dtype], List[List[Tensor]]] \
+        = _group_tensors_by_device_and_dtype([[g.detach() for g in grads]])  # type: ignore[assignment]
+
+    if norm_type == inf:
+        norms = [g.detach().abs().max().to(first_device) for g in grads]
+        total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))
+    else:
+        norms = []
+        for ((device, _), [grads]) in grouped_grads.items():
+            if (foreach is None or foreach) and _has_foreach_support(grads, device=device):
+                norms.extend(torch._foreach_norm(grads, norm_type))
+            elif foreach:
+                raise RuntimeError(f'foreach=True was passed, but can\'t use the foreach API on {device.type} tensors')
+            else:
+                norms.extend([torch.norm(g, norm_type) for g in grads])
+
+        total_norm = torch.norm(torch.stack([norm.to(first_device) for norm in norms]), norm_type)
+        
+    if torch.isnan(total_norm) or torch.isinf(total_norm):
+        for grad in grads:
+            grad.zero_()
+        print('>>>Found nan or inf total_norm, set grads to zero')
+        return torch.tensor(0.)
+
+    if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()):
+        raise RuntimeError(
+            f'The total norm of order {norm_type} for gradients from '
+            '`parameters` is non-finite, so it cannot be clipped. To disable '
+            'this error and scale the gradients by the non-finite norm anyway, '
+            'set `error_if_nonfinite=False`')
+    clip_coef = max_norm / (total_norm + 1e-6)
+    # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so
+    # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization
+    # when the gradients do not reside in CPU memory.
+    clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
+    for ((device, _), [grads]) in grouped_grads.items():
+        if (foreach is None or foreach) and _has_foreach_support(grads, device=device):
+            torch._foreach_mul_(grads, clip_coef_clamped.to(device))  # type: ignore[call-overload]
+        elif foreach:
+            raise RuntimeError(f'foreach=True was passed, but can\'t use the foreach API on {device.type} tensors')
+        else:
+            clip_coef_clamped_device = clip_coef_clamped.to(device)
+            for g in grads:
+                g.detach().mul_(clip_coef_clamped_device)
+
+    return total_norm
+
+def patch_torch_clip_grad_norm():
+    torch.nn.utils.clip_grad_norm_ = clip_grad_norm_

src/utils/resave_model.py

@@ -0,0 +1,70 @@
+import os
+import sys
+sys.path.append(os.getcwd())
+from src.msa.model import MSAForCausalLM
+from src.msa.configuration_msa import MSAConfig
+from src.utils.common import print_model_stats
+from transformers import AutoTokenizer
+
+def save_checkpoint(model, tokenizer, save_model_path):
+    model.save_pretrained(save_model_path)
+    tokenizer.save_pretrained(save_model_path)
+
+def main(origin_model_path, save_model_path):
+    router_layer_idx = os.environ.get("ROUTER_LAYER_IDX", "all")
+    aux_loss = os.environ.get("AUX_LOSS", "false") == "true"
+    lmloss_weigth = float(os.environ.get("LMLOSS_WEIGHT", 1.0))
+    auxloss_weight = float(os.environ.get("AUX_LOSS_WEIGHT", 0.1))
+    recloss_weight = float(os.environ.get("REC_LOSS_WEIGHT", 0.0))
+    ansloss_weight = float(os.environ.get("ANS_LOSS_WEIGHT", 1.0))
+    aux_loss_method = os.environ.get("AUX_LOSS_METHOD", "INFONCE")  # INFONCE, BCE, INFONCE_DECOUPLE, INFONCE_DECOUPLE_FOCAL 
+    decouple_router = os.environ.get("DECOUPLE_ROUTER", "false").lower() == "true"
+    rewrite_position = os.environ.get("REWRITE_POSITION", "false") == "true"
+
+    top_k_docs = int(os.environ.get("TOP_K_DOCS", 2))
+    pooling_kernel_size = int(os.environ.get("POOLING_KERNEL_SIZE", 2))
+
+    head_reduce_method = os.environ.get("HEAD_REDUCE_METHOD", "max")
+    query_reduce_method = os.environ.get("QUERY_REDUCE_METHOD", "max")
+    chunk_reduce_method = os.environ.get("CHUNK_REDUCE_METHOD", "max")
+    decouple_pooling_mode = os.environ.get("DECOUPLE_POOLING_MODE", "mean")
+    infonce_loss_temp = float(os.environ.get("INFONCE_LOSS_TEMP", 0.1))
+
+    msa_config = {
+        "router_layer_idx": router_layer_idx,
+        "aux_loss": aux_loss,
+        "lmloss_weigth": lmloss_weigth,
+        "auxloss_weight": auxloss_weight,
+        "recloss_weight": recloss_weight,
+        "ansloss_weight": ansloss_weight,
+        "aux_loss_method": aux_loss_method,
+        "decouple_router": decouple_router,
+        "rewrite_position": rewrite_position,
+        "top_k_docs": top_k_docs,
+        "pooling_kernel_size": pooling_kernel_size,
+        "infonce_loss_temp": infonce_loss_temp,
+        "head_reduce_method": head_reduce_method,
+        "query_reduce_method": query_reduce_method,
+        "chunk_reduce_method": chunk_reduce_method,
+        "decouple_pooling_mode": decouple_pooling_mode,
+    }
+    # 使用 MSAConfig，它会自动将 msa_config 转换为 DotDict
+    config = MSAConfig.from_pretrained(origin_model_path)
+    config.msa_config = msa_config  # MSAConfig 会自动转换为 DotDict
+    tokenizer = AutoTokenizer.from_pretrained(origin_model_path)
+    model = MSAForCausalLM.from_pretrained(
+        origin_model_path,
+        config=config,
+        torch_dtype="bfloat16",
+    )
+    print_model_stats(model)
+
+    # save
+    save_checkpoint(model, tokenizer, save_model_path)
+
+if __name__ == "__main__":
+    origin_model_path = sys.argv[1]
+    save_model_path = sys.argv[2]
+    main(origin_model_path, save_model_path)
+    print(f"Model has been saved to : {save_model_path}")
+    print("Done")

src/utils/scale.py

@@ -0,0 +1,10 @@
+
+import pickle
+
+def scale_memory(context, query, scale=3):
+    with open(query, "rb") as f:
+        query_metas = pickle.load(f)
+
+    labels_contents = [ref for q_meta  in query_metas for ref in q_meta['reference_list']]
+    copy_reference = [f"{ref}_copy{bias}" for bias, ref in enumerate(context * scale) if ref not in labels_contents]
+    context.extend(copy_reference) 

src/utils/template.py

@@ -0,0 +1,9 @@
+QWEN3_TEMPLATE={
+    "prompt": "<|im_start|>user\n{prompt}<|im_end|>\n",
+    "response": "<|im_start|>assistant\n<think>\n{think_content}\n</think>\n\n{output}<|im_end|>"
+}
+
+QWEN3_INSTRUCT_TEMPLATE={
+    "prompt": "<|im_start|>user\n{prompt}<|im_end|>\n",
+    "response": "<|im_start|>assistant\n{output}<|im_end|>"
+}

src/utils/tools.py

@@ -0,0 +1,120 @@
+import time
+import torch
+import threading
+
+def format_bytes(size_in_bytes):
+    """
+    将字节数转换为人类可读的格式
+    单位为 M 及以下不要小数点，单位为 G 以上保留一位小数，数值不能小于 1
+    """
+    # 定义单位
+    units = ['B', 'K', 'M', 'G', 'T', 'P']
+    
+    # 处理边界情况
+    if size_in_bytes < 1:
+        return "0B"
+    
+    # 计算单位索引
+    unit_index = 0
+    size = float(size_in_bytes)
+    
+    while size >= 1024 and unit_index < len(units) - 1:
+        size /= 1024
+        unit_index += 1
+    
+    # 根据单位决定格式化方式
+    if unit_index <= 2:  # B, K, M 不要小数点
+        if size == int(size):
+            return f"{int(size)}{units[unit_index]}"
+        else:
+            return f"{int(round(size))}{units[unit_index]}"
+    else:  # G 及以上保留一位小数
+        return f"{size:.1f}{units[unit_index]}"
+
+def cumulative_concat(tensors):
+    # 一次性获取所有信息
+    lengths = [len(t) for t in tensors]
+    last_values = [t[-1] for t in tensors]
+    
+    # 计算累积偏移（不包括最后一个tensor）
+    cum_offsets = torch.cumsum(torch.tensor([0] + last_values[:-1]), dim=0)
+    
+    # 构建偏移数组
+    total_length = sum(lengths)
+    offsets = torch.zeros(total_length, dtype=tensors[0].dtype, device=tensors[0].device)
+    
+    # 为每个tensor设置对应的偏移
+    start_idx = 0
+    for i, length in enumerate(lengths):
+        if i > 0:  # 第一个tensor不需要偏移
+            offsets[start_idx:start_idx + length] = cum_offsets[i]
+        start_idx += length
+    
+    # 一次性拼接和添加偏移
+    concatenated = torch.cat(tensors)
+    return concatenated + offsets
+
+    
+class RequestLimiter:
+    def __init__(self, max_concurrent=10):
+        """
+        初始化请求限流器
+        
+        Args:
+            max_concurrent: 最大并发请求数，默认10
+        """
+        self.max_concurrent = max_concurrent
+        self.current_count = 0
+        self.lock = threading.Lock()
+        self.condition = threading.Condition(self.lock)
+    
+    def acquire(self):
+        """
+        获取执行权限，如果超过最大并发数则阻塞
+        
+        Returns:
+            bool: 是否成功获取权限
+        """
+        with self.lock:
+            while self.current_count >= self.max_concurrent:
+                # 等待有请求完成
+                self.condition.wait()
+            self.current_count += 1
+            return True
+    
+    def release(self):
+        """
+        释放一个执行权限，唤醒等待的请求
+        """
+        with self.lock:
+            if self.current_count > 0:
+                self.current_count -= 1
+                # 通知一个等待的线程
+                self.condition.notify()
+
+def compose_input(doc, doc_idx, tokenizer):
+    """组建reference的 input"""
+    new_doc = "<|im_start|>" + f"[{doc_idx}]. {doc}[{doc_idx}]<|im_end|>"
+    return new_doc,tokenizer(new_doc, add_special_tokens=False)
+
+
+class TimePoint:
+    def __init__(self, disabled=False):
+        self.pts = []
+        self.disabled = disabled
+    
+    def add(self, name):
+        if not self.disabled:
+            self.pts.append((name, time.time()))
+
+    def print(self):
+        if len(self.pts) < 2:
+            return
+
+        total = self.pts[-1][1] - self.pts[0][1]
+        s = ""
+        if len(self.pts) > 2:
+            lst1 = self.pts[:-1]
+            lst2 = self.pts[1:]
+            s = " | ".join(f"{item1[0]}->{item2[0]}: {item2[1] - item1[1]:.3f}" for item1, item2 in zip(lst1, lst2))
+        print(f"total: {total:.2f} {s}")