zijie-tian/nano-vllm

Fork 0

Files

Zijie Tian 76af506956 [claudesquad] update from 'multi-request-2' on 13 Jan 26 02:01 CST

2026-01-13 02:01:07 +08:00

11 KiB

Raw Blame History

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):

# 第 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)  # 可能读取旧请求的数据!

场景:

请求 A (32K tokens) 完成，decode_buffer 保留其 KV 数据
请求 B 开始，其 decode_start_pos 可能非零（如果继承了旧状态）
请求 B 在第一个 decode step 时错误地读取了请求 A 的 decode buffer 数据

2.3 潜在问题点

decode_start_pos 计算错误:
- get_decode_start_pos() 使用 id(seq) 作为 key
- Python 对象 ID 可能在请求之间重用
- 如果新 seq 对象的 ID 与旧 seq 相同，可能错误继承旧的 start_pos
decode buffer 残留数据:
- 如果 pos_in_block 在新请求中与旧请求重叠
- get_decode_kv() 会返回旧请求的数据
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):

# 在 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() 末尾):

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

"""最小复现批量模式失败"""
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):

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)

验证命令:

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):

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 末尾):

# 清理 OffloadEngine 状态
if self.offload_engine is not None:
    self.offload_engine.on_sequence_finished()

Phase 4: 实施计划 (pending)

具体修改

文件 1: nanovllm/kvcache/hybrid_manager.py

位置: deallocate() 方法末尾 (第 244 行后)

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

位置: 在类末尾添加新方法

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 (严格限制，不可更改)

# 快速验证 (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%)	最终验收

当前状态

Phase 1: 代码分析
Phase 2: 根本原因分析
Phase 3: Debug 方案设计
Phase 4: 实施计划 ✅ 100/100 PASSED

验证结果

测试	结果	日期
20 样本快速验证	✅ 20/20 (100%)	2026-01-13
100 样本完整验证	✅ 100/100 (100%)	2026-01-13

11 KiB Raw Blame History Unescape Escape