360 lines
11 KiB
Markdown
360 lines
11 KiB
Markdown
# Task Plan: nanovllm CPU Offload 多请求状态污染问题
|
||
|
||
## 问题概述
|
||
|
||
**重要说明**: nanovllm offload 模式目前**不支持 batch**，只能单个 request 顺序执行。问题出在**请求切换**时的状态清理。
|
||
|
||
| 模式 | 测试方式 | 准确率 |
|
||
|------|----------|--------|
|
||
| CPU Offload | 独立进程 (每请求一个进程) | **100%** |
|
||
| CPU Offload | 同进程顺序多请求 | 66% |
|
||
| Non-Offload | 同进程顺序多请求 | 100% |
|
||
|
||
**结论**: 单请求推理正确，问题在于**请求切换**时状态清理不完整。
|
||
|
||
---
|
||
|
||
## Phase 1: 代码分析 (complete)
|
||
|
||
### 1.1 识别状态管理组件
|
||
|
||
**已分析的关键组件**:
|
||
|
||
| 组件 | 文件 | 状态数据 |
|
||
|------|------|----------|
|
||
| `OffloadEngine` | `nanovllm/kvcache/offload_engine.py` | ring buffer, decode buffer, CUDA events |
|
||
| `HybridKVCacheManager` | `nanovllm/kvcache/hybrid_manager.py` | logical blocks, prefilled_blocks, _decode_start_pos, _prefill_len |
|
||
| `LLMEngine` | `nanovllm/engine/llm_engine.py` | generate() 循环，请求生命周期 |
|
||
| `Scheduler` | `nanovllm/engine/scheduler.py` | postprocess() 调用 deallocate() |
|
||
|
||
### 1.2 请求生命周期分析
|
||
|
||
```
|
||
generate()
|
||
→ 多个请求添加到 scheduler
|
||
→ while not finished:
|
||
→ schedule() 获取下一批 seqs
|
||
→ model_runner.run() 执行推理
|
||
→ postprocess() 处理完成的请求
|
||
→ 如果完成: kvcache_manager.deallocate(seq)
|
||
```
|
||
|
||
---
|
||
|
||
## Phase 2: 根本原因分析 (complete)
|
||
|
||
### 2.1 核心问题: OffloadEngine 缺少 reset() 方法
|
||
|
||
**关键发现**: `OffloadEngine` 没有任何重置/清理方法！
|
||
|
||
当请求完成时，`HybridKVCacheManager.deallocate()` 被调用，但它只清理：
|
||
- 逻辑块状态 (`block.reset()`)
|
||
- 物理块引用 (`free_cpu_blocks`, `cpu_block_to_logical`)
|
||
- prefilled_blocks 集合
|
||
- _decode_start_pos / _prefill_len 字典
|
||
|
||
**未被清理的状态** (存在于 OffloadEngine):
|
||
|
||
| 状态 | Shape | 问题 |
|
||
|------|-------|------|
|
||
| `layer_k_cache` | [num_buffers, max_seq_len, kv_heads, head_dim] | 包含旧请求的 KV |
|
||
| `layer_v_cache` | [num_buffers, max_seq_len, kv_heads, head_dim] | 包含旧请求的 KV |
|
||
| `decode_k_buffer` | [num_layers, block_size, kv_heads, head_dim] | 包含旧请求的 decode KV |
|
||
| `decode_v_buffer` | [num_layers, block_size, kv_heads, head_dim] | 包含旧请求的 decode KV |
|
||
|
||
### 2.2 具体污染场景
|
||
|
||
在 `run_layerwise_offload_decode()` (model_runner.py:867-1057):
|
||
|
||
```python
|
||
# 第 969-976 行: 读取之前的 decode KV
|
||
if num_prev_decode_tokens > 0:
|
||
k_decode_prev, v_decode_prev = offload_engine.get_decode_kv(
|
||
layer_id, decode_start_pos, pos_in_block
|
||
)
|
||
ring_k[...].copy_(k_decode_prev) # 可能读取旧请求的数据!
|
||
```
|
||
|
||
**场景**:
|
||
1. 请求 A (32K tokens) 完成，decode_buffer 保留其 KV 数据
|
||
2. 请求 B 开始，其 `decode_start_pos` 可能非零（如果继承了旧状态）
|
||
3. 请求 B 在第一个 decode step 时错误地读取了请求 A 的 decode buffer 数据
|
||
|
||
### 2.3 潜在问题点
|
||
|
||
1. **decode_start_pos 计算错误**:
|
||
- `get_decode_start_pos()` 使用 `id(seq)` 作为 key
|
||
- Python 对象 ID 可能在请求之间重用
|
||
- 如果新 seq 对象的 ID 与旧 seq 相同，可能错误继承旧的 start_pos
|
||
|
||
2. **decode buffer 残留数据**:
|
||
- 如果 `pos_in_block` 在新请求中与旧请求重叠
|
||
- `get_decode_kv()` 会返回旧请求的数据
|
||
|
||
3. **ring buffer 残留数据**:
|
||
- 虽然每次 decode 会从 CPU 加载，但 decode buffer 的数据会被复制过来
|
||
- 如果 decode buffer 有残留，会污染 ring buffer
|
||
|
||
---
|
||
|
||
## Phase 3: Debug 方案设计 (complete)
|
||
|
||
### 3.1 确认的根本原因
|
||
|
||
通过代码分析，确认了两个根本原因：
|
||
|
||
**根本原因 1 (主要)**: `deallocate()` 不调用 `clear_decode_tracking()`
|
||
- 位置: `hybrid_manager.py:218-244`
|
||
- 影响: `_decode_start_pos` 和 `_prefill_len` 字典残留
|
||
- 后果: 如果 `id(seq)` 重用，返回错误的 decode 配置
|
||
|
||
**根本原因 2 (次要)**: decode_buffer 不清理
|
||
- 位置: `offload_engine.py`
|
||
- 影响: `decode_k_buffer/v_buffer` 保留旧 KV
|
||
- 后果: 可能被根本原因 1 触发读取
|
||
|
||
### 3.2 Debug 方案 A: 验证字典残留 (推荐先做)
|
||
|
||
**目标**: 验证 `_decode_start_pos` 字典是否有残留
|
||
|
||
**诊断代码** (添加到 `hybrid_manager.py`):
|
||
```python
|
||
# 在 get_decode_start_pos() 开头添加
|
||
def get_decode_start_pos(self, seq: Sequence) -> int:
|
||
seq_id = id(seq)
|
||
# DEBUG: 检查是否命中旧值
|
||
if seq_id in self._decode_start_pos:
|
||
logger.warning(f"[DEBUG] get_decode_start_pos: CACHE HIT! seq_id={seq_id}, "
|
||
f"cached_value={self._decode_start_pos[seq_id]}, "
|
||
f"expected={(len(seq) - 1) % self._block_size}")
|
||
# ... 原有逻辑
|
||
```
|
||
|
||
**诊断代码** (添加到 `deallocate()` 末尾):
|
||
```python
|
||
def deallocate(self, seq: Sequence) -> None:
|
||
# ... 现有逻辑 ...
|
||
|
||
# DEBUG: 打印未清理的状态
|
||
seq_id = id(seq)
|
||
if seq_id in self._decode_start_pos:
|
||
logger.warning(f"[DEBUG] deallocate: _decode_start_pos NOT CLEARED! "
|
||
f"seq_id={seq_id}, value={self._decode_start_pos[seq_id]}")
|
||
```
|
||
|
||
### 3.3 Debug 方案 B: 最小复现测试
|
||
|
||
**文件**: `tests/test_multi_request_offload_debug.py`
|
||
|
||
```python
|
||
"""最小复现批量模式失败"""
|
||
import os
|
||
import sys
|
||
sys.path.insert(0, os.getcwd())
|
||
|
||
from nanovllm import LLM
|
||
from nanovllm.sampling import SamplingParams
|
||
|
||
# 使用 RULER NIAH 的两个样本
|
||
PROMPTS = [
|
||
# Sample 0 (通常成功)
|
||
"...", # 从 niah_single_1_32k.jsonl 加载
|
||
# Sample 1 (通常失败)
|
||
"...",
|
||
]
|
||
EXPECTED = ["8930103", "4194548"]
|
||
|
||
def main():
|
||
llm = LLM(
|
||
"~/models/Llama-3.1-8B-Instruct",
|
||
max_model_len=33792,
|
||
max_num_batched_tokens=33792,
|
||
enable_cpu_offload=True,
|
||
num_gpu_blocks=4,
|
||
kvcache_block_size=1024,
|
||
enforce_eager=True,
|
||
)
|
||
|
||
params = SamplingParams(temperature=0.1, max_tokens=50)
|
||
|
||
# 连续处理两个请求
|
||
for i, (prompt, expected) in enumerate(zip(PROMPTS, EXPECTED)):
|
||
print(f"\n{'='*60}")
|
||
print(f"Sample {i}: Expected = {expected}")
|
||
|
||
# 打印关键状态
|
||
kvm = llm.model_runner.kvcache_manager
|
||
print(f" _decode_start_pos 字典大小: {len(kvm._decode_start_pos)}")
|
||
print(f" _prefill_len 字典大小: {len(kvm._prefill_len)}")
|
||
|
||
outputs = llm.generate([prompt], params, use_tqdm=False)
|
||
output_text = outputs[0]["text"]
|
||
|
||
passed = expected in output_text
|
||
print(f" Output: {output_text[:100]}...")
|
||
print(f" Status: {'PASS' if passed else 'FAIL'}")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
```
|
||
|
||
### 3.4 Debug 方案 C: 快速修复验证
|
||
|
||
**目标**: 验证修复 `deallocate()` 是否解决问题
|
||
|
||
**修改** (`hybrid_manager.py:218-244`):
|
||
```python
|
||
def deallocate(self, seq: Sequence) -> None:
|
||
"""Release all blocks for a sequence."""
|
||
for logical_id in reversed(seq.block_table):
|
||
# ... 现有逻辑 ...
|
||
|
||
seq.num_cached_tokens = 0
|
||
seq.block_table.clear()
|
||
|
||
# === 新增: 清理 decode tracking ===
|
||
self.clear_decode_tracking(seq)
|
||
```
|
||
|
||
**验证命令**:
|
||
```bash
|
||
CUDA_VISIBLE_DEVICES=0 PYTHONPATH=.:$PYTHONPATH python tests/test_ruler_niah.py \
|
||
--model ~/models/Llama-3.1-8B-Instruct \
|
||
--enable-offload \
|
||
--sample-indices 0,1,2,3,4 \
|
||
--verbose
|
||
```
|
||
|
||
### 3.5 Debug 方案 D: 添加 OffloadEngine 清理 (防御性)
|
||
|
||
**目标**: 进一步隔离请求状态
|
||
|
||
**添加方法** (`offload_engine.py`):
|
||
```python
|
||
def on_sequence_finished(self):
|
||
"""清理请求完成后的状态"""
|
||
# 清零 decode buffer (防止残留数据被读取)
|
||
self.decode_k_buffer.zero_()
|
||
self.decode_v_buffer.zero_()
|
||
logger.debug("OffloadEngine: decode buffer cleared")
|
||
```
|
||
|
||
**调用点** (`hybrid_manager.py:deallocate` 末尾):
|
||
```python
|
||
# 清理 OffloadEngine 状态
|
||
if self.offload_engine is not None:
|
||
self.offload_engine.on_sequence_finished()
|
||
```
|
||
|
||
---
|
||
|
||
## Phase 4: 实施计划 (pending)
|
||
|
||
### 推荐执行顺序
|
||
|
||
1. **Step 4.1**: 实施修复
|
||
- 修改 `hybrid_manager.py:deallocate()` 添加 `clear_decode_tracking(seq)`
|
||
|
||
2. **Step 4.2**: 快速验证 (20 样本连续执行)
|
||
- **一次调用** `test_ruler_niah.py`，连续执行 20 个样本
|
||
- **不重启框架**，验证请求切换是否正确
|
||
- 目标: 20/20 全部通过
|
||
|
||
3. **Step 4.3**: 完整验证 (100 样本)
|
||
- 运行 100 个样本的 RULER NIAH 测试
|
||
- 目标: 100/100 全部通过 (准确率从 66% → 100%)
|
||
|
||
4. **Step 4.4**: 防御性修复 (可选)
|
||
- 添加 `OffloadEngine.on_sequence_finished()` 方法
|
||
- 清零 decode buffer 作为额外保险
|
||
|
||
### 具体修改
|
||
|
||
**文件 1**: `nanovllm/kvcache/hybrid_manager.py`
|
||
|
||
位置: `deallocate()` 方法末尾 (第 244 行后)
|
||
|
||
```python
|
||
def deallocate(self, seq: Sequence) -> None:
|
||
"""Release all blocks for a sequence."""
|
||
for logical_id in reversed(seq.block_table):
|
||
# ... 现有逻辑 (218-242 行) ...
|
||
|
||
seq.num_cached_tokens = 0
|
||
seq.block_table.clear()
|
||
|
||
# ============ 新增: 清理 decode tracking ============
|
||
self.clear_decode_tracking(seq)
|
||
```
|
||
|
||
**文件 2** (可选): `nanovllm/kvcache/offload_engine.py`
|
||
|
||
位置: 在类末尾添加新方法
|
||
|
||
```python
|
||
def on_sequence_finished(self):
|
||
"""清理请求完成后的状态 (防御性清理)"""
|
||
self.decode_k_buffer.zero_()
|
||
self.decode_v_buffer.zero_()
|
||
```
|
||
|
||
---
|
||
|
||
## 关键文件清单
|
||
|
||
| 文件 | 相关行号 | 说明 |
|
||
|------|----------|------|
|
||
| `nanovllm/kvcache/hybrid_manager.py` | 218-244 | `deallocate()` - **需要修改** |
|
||
| `nanovllm/kvcache/hybrid_manager.py` | 538-549 | `clear_decode_tracking()` - 已存在 |
|
||
| `nanovllm/kvcache/hybrid_manager.py` | 485-505 | `get_decode_start_pos()` - 问题读取点 |
|
||
| `nanovllm/kvcache/hybrid_manager.py` | 519-537 | `get_prefill_len()` - 问题读取点 |
|
||
| `nanovllm/kvcache/offload_engine.py` | 40-145 | `__init__` - 状态初始化 |
|
||
| `nanovllm/kvcache/offload_engine.py` | (新增) | `on_sequence_finished()` - 可选防御 |
|
||
| `nanovllm/engine/model_runner.py` | 867-1057 | `run_layerwise_offload_decode()` |
|
||
| `nanovllm/engine/model_runner.py` | 969-976 | decode buffer 读取 (污染点) |
|
||
|
||
---
|
||
|
||
## 验证命令
|
||
|
||
**指定 GPU: 1** (严格限制，不可更改)
|
||
|
||
```bash
|
||
# 快速验证 (20 样本连续执行，不重启框架)
|
||
# 目标: 20/20 通过
|
||
CUDA_VISIBLE_DEVICES=1 PYTHONPATH=.:$PYTHONPATH python tests/test_ruler_niah.py \
|
||
--model ~/models/Llama-3.1-8B-Instruct \
|
||
--enable-offload \
|
||
--sample-indices 0-19 \
|
||
--verbose
|
||
|
||
# 完整验证 (100 样本)
|
||
# 目标: 100/100 通过 (最终验收)
|
||
CUDA_VISIBLE_DEVICES=1 PYTHONPATH=.:$PYTHONPATH python tests/test_ruler_niah.py \
|
||
--model ~/models/Llama-3.1-8B-Instruct \
|
||
--enable-offload \
|
||
--quiet
|
||
```
|
||
|
||
**验收标准**:
|
||
| 测试 | 样本数 | 通过要求 | 说明 |
|
||
|------|--------|----------|------|
|
||
| 快速验证 | 20 | 20/20 (100%) | 一次调用，连续执行，验证请求切换 |
|
||
| 完整验证 | 100 | 100/100 (100%) | 最终验收 |
|
||
|
||
---
|
||
|
||
## 当前状态
|
||
|
||
- [x] Phase 1: 代码分析
|
||
- [x] Phase 2: 根本原因分析
|
||
- [x] Phase 3: Debug 方案设计
|
||
- [x] Phase 4: 实施计划 ✅ 100/100 PASSED
|
||
|
||
### 验证结果
|
||
|
||
| 测试 | 结果 | 日期 |
|
||
|------|------|------|
|
||
| 20 样本快速验证 | ✅ 20/20 (100%) | 2026-01-13 |
|
||
| 100 样本完整验证 | ✅ 100/100 (100%) | 2026-01-13 |
|