nano-vllm/docs/chunked_prefill_integration_plan.md

Chunked Prefill 集成计划

目标: 将 tzj/minference 分支的 chunked prefill 机制移植到 tzj/vs_offload 分支

创建日期: 2026-01-18 基础分支: tzj/vs_offload 源分支: tzj/minference

---

目标

在 tzj/vs_offload 分支上实现 chunked prefill + layerwise offload 机制，支持在 24GB RTX 3090 上运行任意长度的推理（4M, 8M, 16M+ tokens）。

---

核心问题

tzj/vs_offload 分支的局限性

当前 tzj/vs_offload 分支的 GPU ring buffer 按 max_seq_len 分配，导致 GPU 内存随序列长度线性增长：

# 当前设计
self.layer_k_cache = torch.zeros(
    num_kv_buffers,      # e.g., 4
    max_seq_len,         # e.g., 131072 tokens
    kv_heads,
    head_dim,
    dtype=dtype, device="cuda"
)

问题：

• GPU 内存需求 ~ max_seq_len × 4 × 8 × 128 × 2 bytes
• 对于超长序列不可行：
  • 4M tokens → ~64 GB GPU 内存 ❌
  • 8M tokens → ~128 GB GPU 内存 ❌

解决方案：Block-Based 设计

tzj/minference 分支采用 block-based 设计，GPU 内存固定：

# Block-based 设计
self.k_cache_gpu = torch.zeros(
    num_gpu_blocks,      # e.g., 2
    block_size,          # e.g., 1024 tokens (固定!)
    kv_heads,
    head_dim,
    dtype=dtype, device="cuda"
)
# GPU 内存: ~4 MB (固定，不随序列长度增长)

优势：

• GPU 内存固定（~1.6 GB），不随序列长度增长
• 24GB RTX 3090 可运行 4M+ tokens
• 通过 chunked prefill 分块处理超长序列

---

内存布局对比

组件	tzj/vs_offload	tzj/minference	说明
GPU Ring Buffer	[num_kv_buffers, max_seq_len, ...]	[num_gpu_blocks, block_size, ...]	minference 无 layer 维度
GPU 内存	~2.15 GB (128K) → ~64 GB (4M)	~4 MB (固定)	minference 节省显著
Prefill Buffer	❌ 无	✅ [num_layers, block_size, ...]	minference 独有
Pipeline Buffers	❌ 无	✅ 双缓冲区 [blocks, block_size, ...]	minference 独有
CPU Cache	[num_layers, num_cpu_blocks, block_size, ...]	相同	一致

序列长度支持对比

序列长度	vs_offload GPU 内存	minference GPU 内存	RTX 3090 (24GB)
128K tokens	~2.15 GB	~4 MB	✅ 两者均可
1M tokens	~16 GB	~4 MB	✅ 两者均可
4M tokens	~64 GB ❌	~4 MB ✅	仅 minference 可行
8M tokens	~128 GB ❌	~4 MB ✅	仅 minference 可行
16M+ tokens	~256 GB+ ❌	~4 MB ✅	仅 minference 可行

---

关键设计原则

1. Block-Based 设计：按 block_size (1024 tokens) 组织，支持 chunked prefill
2. GPU 内存固定：不随序列长度增长，是 constant factor
3. CPU 内存线性缩放：num_cpu_blocks = ceil(seq_len / block_size)
4. Unified Ring Buffer：无 layer 维度，所有层共享 slots
5. 完全并行 offload：per-layer buffer 最大化 PCIe 带宽

---

统一内存布局设计

GPU Memory Layout

class OffloadEngine:
    # 1. Unified Ring Buffer - Block-based，无 layer 维度
    self.k_cache_gpu = torch.zeros(
        num_gpu_blocks,  # e.g., 2
        block_size,      # e.g., 1024
        kv_heads,
        head_dim,
        dtype=dtype, device="cuda"
    )  # ~4 MB (固定)

    # 2. Per-layer Prefill Buffer - 完全并行 offload
    self.prefill_k_buffer = torch.zeros(
        num_layers, block_size, kv_heads, head_dim,
        dtype=dtype, device="cuda"
    )  # ~58 MB (固定)

    # 3. Cross-layer Pipeline Buffers - Double-buffering
    self.layer_k_buffer_a = torch.zeros(
        max_prefill_blocks, block_size, kv_heads, head_dim,
        dtype=dtype, device="cuda"
    )  # ~512 MB (固定)
    self.layer_k_buffer_b = torch.zeros(...)  # ~512 MB (固定)

    # 4. Per-layer Decode Buffer
    self.decode_k_buffer = torch.zeros(
        num_layers, block_size, kv_heads, head_dim,
        dtype=dtype, device="cuda"
    )  # ~58 MB (固定)

    # GPU 总计：~1.6 GB (固定，不随序列长度增长)

CPU Memory Layout

    # CPU Cache - 有 block 维度
    self.k_cache_cpu = torch.zeros(
        num_layers,
        num_cpu_blocks,  # 随序列长度缩放
        block_size,
        kv_heads,
        head_dim,
        dtype=dtype, device="cpu", pin_memory=True
    )
    # 128K tokens: ~2.9 GB
    # 1M tokens: ~5.8 GB
    # 4M tokens: ~23.3 GB

---

Chunked Prefill 流程

Prefill 阶段

For each chunk:
├── 1. Prepare chunk input (block_size tokens)
├── 2. Get ring buffer slot: slot = chunk_idx % num_gpu_blocks
├── 3. Load previous KV chunks to ring slots[1..N-1]
├── 4. Model Forward (all layers)
│   For each layer:
│   ├── Load previous KV from ring slots
│   ├── Compute attention (current chunk + previous)
│   ├── Write KV to prefill_buffer[layer_id]  ← Per-layer!
│   └── Async offload to CPU (parallel across layers)
├── 5. Merge attention outputs (LSE)
└── 6. Record compute done for slot

Key: Per-layer prefill buffer → Layer 0 offload || Layer 1 compute || Layer 2 load ...

Decode 阶段

├── 1. Setup pipeline: preload Layer 0 to buffer_a
├── 2. For each layer:
│   ├── Get KV from pipeline buffer (a or b)
│   ├── Trigger preload of next layer to other buffer
│   ├── Compute attention
│   └── Store to decode buffer
└── 3. End pipeline

Key: Double-buffering → Layer N compute || Layer N+1 load

---

合并策略

基础分支选择：tzj/vs_offload

原因：

1. 更完善的文档系统
2. 更完整的 sparse attention 实现（QUEST, XAttention 等）
3. 更清晰的代码组织和注释
4. 更活跃的开发维护

移植策略

从 tzj/minference 移植：

1. GPU cache 内存布局（无 layer 维度，block-based）
2. Per-layer prefill buffer
3. Cross-layer pipeline buffers
4. Chunked prefill 流程
5. LSE 在线合并机制

保留 tzj/vs_offload 优势：

1. 文档系统
2. Sparse policy 架构
3. 代码组织和注释

---

Sparse Policy 策略

策略：保留架构，现阶段仅实现 FULL

• 保留 sparse policy 的架构设计和接口
• 预留 扩展接口给未来的 QUEST 等其他策略
• 现阶段仅实现 FULL 策略，确保正确性和稳定性

实现

class SparsePolicy(ABC):
    @property
    def supports_prefill(self) -> bool:
        return False

    @property
    def supports_decode(self) -> bool:
        return True

    def on_prefill_offload(self, cpu_block_id, layer_id, k_cache, num_valid_tokens):
        """预留给未来策略（如 QUEST 收集元数据）"""
        pass

    def select_blocks(self, available_blocks, context) -> List[int]:
        """FULL: 返回所有可用块"""
        return available_blocks

class FullAttentionPolicy(SparsePolicy):
    @property
    def supports_prefill(self) -> bool:
        return True

    @property
    def supports_decode(self) -> bool:
        return True

---

关键 API

Ring Buffer 管理

# Prefill 阶段
get_write_slot_for_prefill(chunk_idx) -> slot_idx
get_load_slots_for_prefill(write_slot_idx) -> [slot_ids]

# Decode 阶段
get_load_slots_for_decode() -> [slot_ids] (excludes decode_slot)

Per-layer 操作

# 加载
load_to_slot_layer(slot_idx, layer_id, cpu_block_id)
wait_slot_layer(slot_idx)

# Prefill buffer
get_prefill_buffer(layer_id) -> (k, v)
offload_prefill_buffer_async(layer_id, cpu_block_id, num_tokens)
wait_prefill_offload(layer_id)

# Pipeline
start_decode_pipeline(cpu_block_ids)
get_decode_layer_kv(layer_id, num_blocks) -> (k, v)
end_decode_pipeline()

---

实施阶段

Phase 1: 内存布局重构

• 修改 GPU cache 为 unified ring buffer
• 添加 per-layer prefill buffer
• 添加 cross-layer pipeline buffers

Phase 2: API 实现

• 实现 ring buffer slot 管理 API
• 实现 per-layer prefill offload API
• 实现 cross-layer pipeline API

Phase 3: 集成到 Attention Layer

• 修改 attention forward 流程
• 集成 per-layer prefill buffer
• 集成 cross-layer pipeline

Phase 4: 集成到 Model Runner

• 实现 chunked prefill 流程
• 集成 LSE 合并
• 优化流水线

Phase 5: Sparse Policy 集成（FULL）

• 设计统一的策略接口
• 实现 FullAttentionPolicy
• 预留 QUEST 等未来策略的扩展接口

---

关键决策

1. Block-Based 设计优先：支持任意长度推理的核心
2. 采用 tzj/minference 的内存布局：GPU cache 无 layer 维度 + block-based
3. 以 tzj/vs_offload 为基础分支：更好的文档和代码组织
4. 分阶段合并策略：降低复杂度，便于验证
5. Sparse Policy - FULL 优先：保留架构，现阶段仅实现 FULL

---

预期结果

内存使用（28层模型，block_size=1024）

组件	内存
GPU Unified Ring Buffer	~4 MB
GPU Per-layer Prefill Buffer	~58 MB
GPU Pipeline Buffers (×2)	~1 GB
GPU Decode Buffer	~58 MB
GPU 总计	~1.6 GB (固定)
CPU Cache (4M tokens)	~23.3 GB
总计 (4M tokens)	~24.9 GB ✅ 适配 24GB RTX 3090

性能支持

• ✅ 支持 4M, 8M, 16M+ tokens 的推理
• ✅ GPU 内存固定，不随序列长度增长
• ✅ 完全并行的 layerwise offload
• ✅ Cross-layer 流水线优化

---

参考

• OffloadEngine: nanovllm/kvcache/offload_engine.py
• Attention Layer: nanovllm/layers/attention.py
• Model Runner: nanovllm/engine/model_runner.py
• Sparse Policy: nanovllm/kvcache/sparse/policy.py