[claudesquad] update from 'int-offload-1' on 08 Jan 26 19:44 CST

2026-01-08 19:44:29 +08:00
parent b5c0ef3b7a
commit 85bcca3d17
2 changed files with 604 additions and 0 deletions
--- a/task_plan.md
+++ b/task_plan.md
@@ -0,0 +1,399 @@
+# Task Plan: Layerwise Offload Refactoring
+
+## Goal
+Refactor layerwise offload to use proper OffloadEngine API, pre-allocate buffers, remove chunked prefill code, and pass needle test.
+
+## Phases
+- [ ] Phase 1: Add layerwise API to OffloadEngine
+- [ ] Phase 2: Pre-allocate buffers in ModelRunner
+- [ ] Phase 3: Refactor run_layerwise_offload_prefill()
+- [ ] Phase 4: Refactor run_layerwise_offload_decode()
+- [ ] Phase 5: Remove chunked prefill code
+- [ ] Phase 6: Verify with needle test
+
+## Key Questions
+1. Should we keep chunked_attention.py for MInference use?
+2. What's the max_seq_len for buffer pre-allocation?
+3. Should we implement incremental refactoring or all at once?
+
+## Decisions Made
+- Use FullAttentionPolicy for initial testing (per user request)
+- Focus on correctness first, then optimize async overlap
+- **GPU KV Cache使用Ring Buffer策略** (用户建议):
+  - 使用N个buffer (可配置，默认4个) 形成ring buffer
+  - 比固定2个buffer更灵活，流水线深度更深
+  - 可以预加载多层，更好地隐藏H2D延迟
+  - 例如: buffer[i] compute, buffer[(i+1)%N] load, buffer[(i+2)%N] load...
+
+## Errors Encountered
+(none yet)
+
+## Status
+**Currently in Phase 0** - Planning complete, awaiting user approval
+
+---
+
+## Detailed Implementation Plan
+
+### Phase 1: Modify OffloadEngine GPU Memory Layout + Add Layerwise API
+
+**File**: `nanovllm/kvcache/offload_engine.py`
+
+#### 1.1 新的GPU内存布局 (Ring Buffer)
+
+**设计原则**:
+- 不追求极致的peak memory优化，而是保证流水线正确性和性能
+- Ring buffer层数可从外部配置 (通过config或参数)
+- 默认4层，可以根据GPU内存和H2D带宽调整
+
+```python
+# ========== Ring-Buffered GPU KV Cache for Layerwise Offload ==========
+#
+# 参数: num_kv_buffers (外部可配置，默认4)
+#
+# Ring Buffer流水线 (以4个buffer为例):
+#   Buffer 0: [Load L0] → [Compute L0] ──────────────────────────► [Load L4]
+#   Buffer 1:           [Load L1] → [Compute L1] ──────────────────────────►
+#   Buffer 2:                     [Load L2] → [Compute L2] ────────────────►
+#   Buffer 3:                               [Load L3] → [Compute L3] ──────►
+#
+# 优势:
+# - 流水线深度 = num_kv_buffers - 1
+# - 可以预加载多层，更好地隐藏H2D延迟
+# - 比固定2层更灵活
+
+def __init__(
+    self,
+    ...,
+    num_kv_buffers: int = 4,  # 外部可配置的ring buffer层数
+):
+    self.num_kv_buffers = num_kv_buffers
+
+    # Shape: [num_kv_buffers, max_seq_tokens, kv_heads, head_dim]
+    self.layer_k_cache = torch.zeros(
+        num_kv_buffers, max_seq_tokens, num_kv_heads, head_dim,
+        dtype=dtype, device="cuda"
+    )
+    self.layer_v_cache = torch.zeros(
+        num_kv_buffers, max_seq_tokens, num_kv_heads, head_dim,
+        dtype=dtype, device="cuda"
+    )
+
+    # Per-buffer events for H2D completion
+    self.buffer_load_events = [torch.cuda.Event() for _ in range(num_kv_buffers)]
+
+# 内存开销计算 (Qwen3-4B, 128K tokens):
+# - kv_heads=8, head_dim=128, dtype=bf16
+# - 单层: 128K × 8 × 128 × 2 = 256 MB
+# - 4层ring buffer: 4 × 256 MB = 1 GB
+# - 对比28层全部在GPU: 28 × 256 MB = 7.2 GB
+# - **节省**: 7.2 GB - 1 GB = 6.2 GB
+```
+
+**配置传递路径**:
+```
+LLM(num_kv_buffers=4)
+  → Config.num_kv_buffers
+    → OffloadEngine(num_kv_buffers=config.num_kv_buffers)
+```
+
+**移除旧的ring buffer设计**:
+```python
+# 移除: k_cache_gpu, v_cache_gpu (chunked prefill用的ring buffer)
+# 移除: ring_slot_ready, ring_slot_offload_done, ring_slot_compute_done
+# 移除: slot_transfer_streams
+# 保留: prefill_offload_streams (用于D2H), compute_stream
+```
+
+#### 1.2 新的Layerwise API方法
+
+```python
+# ========== Prefill: Async D2H Offload ==========
+def offload_layer_kv_async(
+    self, layer_id: int, k: Tensor, v: Tensor,
+    cpu_block_ids: list[int], total_tokens: int
+) -> None:
+    """Async offload layer KV to CPU using per-layer stream."""
+    stream = self.prefill_offload_streams[layer_id]
+    with torch.cuda.stream(stream):
+        stream.wait_stream(self.compute_stream)  # Wait for compute
+        for i, cpu_block_id in enumerate(cpu_block_ids):
+            start = i * self.block_size
+            end = min(start + self.block_size, total_tokens)
+            self.k_cache_cpu[layer_id, cpu_block_id, :end-start].copy_(
+                k[start:end], non_blocking=True
+            )
+            self.v_cache_cpu[layer_id, cpu_block_id, :end-start].copy_(
+                v[start:end], non_blocking=True
+            )
+        self.prefill_offload_events[layer_id].record(stream)
+
+def wait_layer_offload(self, layer_id: int) -> None:
+    """Wait for specific layer's offload to complete."""
+    self.compute_stream.wait_event(self.prefill_offload_events[layer_id])
+
+# ========== Decode: Ring-Buffered H2D Load ==========
+def load_layer_kv_to_buffer(
+    self, buffer_idx: int, layer_id: int,
+    cpu_block_ids: list[int], valid_tokens_per_block: list[int]
+) -> None:
+    """
+    Async load layer KV from CPU to specified ring buffer slot.
+
+    Args:
+        buffer_idx: Ring buffer slot index (0 to num_kv_buffers-1)
+        layer_id: Which layer's KV to load
+        cpu_block_ids: CPU block IDs containing this layer's KV
+        valid_tokens_per_block: Number of valid tokens in each block
+    """
+    stream = self.layer_load_streams[buffer_idx]  # 每个buffer有独立的stream
+    with torch.cuda.stream(stream):
+        # 等待该buffer上一次compute完成 (防止覆盖正在使用的数据)
+        stream.wait_event(self.buffer_compute_done_events[buffer_idx])
+
+        offset = 0
+        for i, cpu_block_id in enumerate(cpu_block_ids):
+            valid_tokens = valid_tokens_per_block[i]
+            self.layer_k_cache[buffer_idx, offset:offset+valid_tokens].copy_(
+                self.k_cache_cpu[layer_id, cpu_block_id, :valid_tokens],
+                non_blocking=True
+            )
+            self.layer_v_cache[buffer_idx, offset:offset+valid_tokens].copy_(
+                self.v_cache_cpu[layer_id, cpu_block_id, :valid_tokens],
+                non_blocking=True
+            )
+            offset += valid_tokens
+        self.buffer_load_events[buffer_idx].record(stream)
+
+def wait_buffer_load(self, buffer_idx: int) -> None:
+    """Wait for buffer load to complete on compute_stream."""
+    self.compute_stream.wait_event(self.buffer_load_events[buffer_idx])
+
+def get_buffer_kv(self, buffer_idx: int, total_tokens: int) -> tuple[Tensor, Tensor]:
+    """Get KV from specified ring buffer slot."""
+    return (
+        self.layer_k_cache[buffer_idx, :total_tokens],
+        self.layer_v_cache[buffer_idx, :total_tokens]
+    )
+
+def record_buffer_compute_done(self, buffer_idx: int) -> None:
+    """Record that compute on this buffer is done (allows next load to reuse it)."""
+    self.buffer_compute_done_events[buffer_idx].record(self.compute_stream)
+```
+
+#### 1.3 Ring Buffer所需的额外资源
+
+```python
+# Per-buffer streams (并行加载多个buffer)
+self.layer_load_streams = [torch.cuda.Stream() for _ in range(num_kv_buffers)]
+
+# Per-buffer events
+self.buffer_load_events = [torch.cuda.Event() for _ in range(num_kv_buffers)]
+self.buffer_compute_done_events = [torch.cuda.Event() for _ in range(num_kv_buffers)]
+
+# 初始化: 标记所有buffer为"compute done" (允许首次加载)
+for event in self.buffer_compute_done_events:
+    event.record()
+```
+
+### Phase 2: Pre-allocate Buffers in ModelRunner
+
+**File**: `nanovllm/engine/model_runner.py`
+
+Add in `__init__()`:
+```python
+def _allocate_layerwise_buffers(self):
+    max_seq_len = self.config.max_model_len
+    hidden_size = self.config.hf_config.hidden_size
+    num_heads = self.config.hf_config.num_attention_heads
+    num_kv_heads = self.config.hf_config.num_key_value_heads
+    head_dim = hidden_size // num_heads
+
+    # QKV buffer for prefill
+    self.prefill_qkv_buffer = torch.empty(
+        max_seq_len, hidden_size + 2 * num_kv_heads * head_dim,
+        dtype=self.dtype, device="cuda"
+    )
+
+    # Decode buffers (single token)
+    self.decode_qkv_buffer = torch.empty(
+        1, hidden_size + 2 * num_kv_heads * head_dim,
+        dtype=self.dtype, device="cuda"
+    )
+```
+
+### Phase 3: Refactor run_layerwise_offload_prefill()
+
+**Key changes**:
+1. Use `offload_engine.compute_stream` for all computation
+2. Use `offload_layer_kv_async()` instead of `_offload_layer_kv_to_cpu_sync()`
+3. Enable overlap: layer N offload overlaps with layer N+1 compute
+4. Remove `torch.cuda.synchronize()`
+
+```python
+def run_layerwise_offload_prefill(self, seqs):
+    offload_engine = self.kvcache_manager.offload_engine
+    compute_stream = offload_engine.compute_stream
+
+    with torch.cuda.stream(compute_stream):
+        for layer_id in range(num_layers):
+            # Wait for previous layer's offload buffer to be safe
+            if layer_id > 0:
+                offload_engine.wait_layer_offload(layer_id - 1)
+
+            # Compute (using pre-allocated buffers where possible)
+            q, k, v = compute_layer_qkv(...)
+            attn_out = flash_attn_varlen_func(q, k, v, causal=True)
+            hidden_states = compute_mlp(...)
+
+            # Async offload (overlaps with next layer)
+            offload_engine.offload_layer_kv_async(layer_id, k, v, cpu_block_ids, total_tokens)
+
+    # Wait for final layer
+    offload_engine.wait_layer_offload(num_layers - 1)
+```
+
+### Phase 4: Refactor run_layerwise_offload_decode()
+
+**Key changes**:
+1. 使用Ring Buffer实现compute/transfer overlap
+2. N个buffer循环使用 (N = num_kv_buffers, 外部可配置)
+3. 使用stream events而非global sync
+4. 流水线深度 = N-1 (可预加载N-1层)
+
+**Ring Buffer流水线示意** (以4个buffer为例):
+```
+时间 ────────────────────────────────────────────────────────────────────────►
+
+Buffer 0: [Load L0] ─► [Compute L0] ────────────────────────► [Load L4] ─►
+Buffer 1:            [Load L1] ─► [Compute L1] ────────────────────────►
+Buffer 2:                       [Load L2] ─► [Compute L2] ────────────────►
+Buffer 3:                                  [Load L3] ─► [Compute L3] ────►
+
+流水线深度 = 3 (同时预加载3层)
+```
+
+```python
+def run_layerwise_offload_decode(self, seqs):
+    offload_engine = self.kvcache_manager.offload_engine
+    compute_stream = offload_engine.compute_stream
+    num_buffers = offload_engine.num_kv_buffers
+
+    # 计算每个block的valid tokens
+    valid_tokens_per_block = self._compute_valid_tokens(cpu_block_table, total_prefill_tokens)
+
+    # Phase 1: 预加载前N层到ring buffer (填满流水线)
+    num_preload = min(num_buffers, num_layers)
+    for i in range(num_preload):
+        offload_engine.load_layer_kv_to_buffer(
+            i, i, cpu_block_table, valid_tokens_per_block
+        )
+
+    # Phase 2: 主循环 - compute当前层，load下一层
+    with torch.cuda.stream(compute_stream):
+        for layer_id in range(num_layers):
+            # 1. 计算当前buffer index (ring)
+            current_buffer = layer_id % num_buffers
+
+            # 2. 等待当前buffer的加载完成
+            offload_engine.wait_buffer_load(current_buffer)
+
+            # 3. 开始加载下一层到同一buffer (buffer被复用)
+            #    下一层 = layer_id + num_buffers (因为当前层用完后buffer可复用)
+            next_layer_to_load = layer_id + num_buffers
+            if next_layer_to_load < num_layers:
+                offload_engine.load_layer_kv_to_buffer(
+                    current_buffer, next_layer_to_load, cpu_block_table, valid_tokens_per_block
+                )
+
+            # 4. 获取当前buffer的KV并计算
+            k_prefill, v_prefill = offload_engine.get_buffer_kv(current_buffer, total_prefill_tokens)
+
+            # 5. 计算新token的QKV
+            q_new, k_new, v_new = self._compute_decode_qkv(layer_id, hidden_states)
+
+            # 6. 拼接并计算attention
+            k_full = torch.cat([k_prefill, k_decode_prev, k_new], dim=0)
+            v_full = torch.cat([v_prefill, v_decode_prev, v_new], dim=0)
+            attn_out = flash_attn_varlen_func(q_new, k_full, v_full, causal=False)
+
+            # 7. 标记当前buffer的compute完成 (允许后续load复用这个buffer)
+            offload_engine.record_buffer_compute_done(current_buffer)
+
+            # 8. 存储新KV到decode buffer
+            offload_engine.decode_k_buffer[layer_id, pos].copy_(k_new.squeeze(0))
+            offload_engine.decode_v_buffer[layer_id, pos].copy_(v_new.squeeze(0))
+
+            # 9. MLP
+            hidden_states = self._compute_mlp(layer_id, attn_out)
+
+    # Block满时offload (使用async API)
+    if block_is_full:
+        offload_engine.offload_decode_buffer_async(cpu_block_id)
+        # 注意: 这里不需要立即wait，可以在下一个decode step开始前wait
+```
+
+**优势**:
+- Compute和H2D transfer完全overlap
+- 流水线深度可配置 (num_kv_buffers-1)
+- 没有global `torch.cuda.synchronize()`
+- 使用stream events进行细粒度同步
+- Buffer在layer_id + num_buffers时自动复用
+
+### Phase 5: Remove Chunked Prefill Code
+
+**Files to modify**:
+
+| File | Remove |
+|------|--------|
+| `nanovllm/layers/attention.py` | `_chunked_prefill_attention()`, `_chunked_decode_attention()`, `_sync_load_previous_chunks()`, `_ring_buffer_pipeline_load()`, `_decode_ring_buffer_pipeline()`, `_decode_with_layer_pipeline()` |
+| `nanovllm/utils/context.py` | `is_chunked_prefill`, `prev_kv_ranges`, `chunk_offset`, `chunked_seq`, `decode_pos_in_block`, `decode_start_pos_in_block`, `current_chunk_idx` |
+| `nanovllm/kvcache/chunked_attention.py` | Keep for MInference (or remove if unused) |
+
+Simplify `Attention.forward()` to:
+```python
+def forward(self, q, k, v):
+    if context.is_prefill:
+        if context.sparse_prefill_policy:
+            return policy.sparse_prefill_attention(q, k, v, self.layer_id)
+        else:
+            return flash_attn_varlen_func(q, k, v, causal=True)
+    else:
+        return flash_attn_with_kvcache(q, k_cache, v_cache, causal=True)
+```
+
+### Phase 6: Verification
+
+**Test command**:
+```bash
+PYTHONPATH=/home/zijie/.claude-squad/worktrees/zijie/int-offload-1_188890c8699249f7:$PYTHONPATH \
+python tests/test_needle.py \
+    --model ~/models/Qwen3-4B-Instruct-2507/ \
+    --max-model-len 32768 \
+    --input-len 8192 \
+    --enable-offload \
+    --block-size 1024 \
+    --num-gpu-blocks 2
+```
+
+**Success criteria**: `test_needle: PASSED`
+
+---
+
+## Current Issues Summary
+
+| Issue | Location | Solution |
+|-------|----------|----------|
+| Direct `.copy_()` bypassing OffloadEngine | `model_runner.py:798-804` | Use `offload_layer_kv_async()` |
+| `torch.cuda.synchronize()` | `model_runner.py:804` | Use stream events |
+| Intermediate memory not pre-allocated | `model_runner.py:508-517` | Pre-allocate in `__init__()` |
+| Chunked prefill code unused | `attention.py`, `context.py` | Remove entirely |
+
+---
+
+## Critical Files
+
+- `nanovllm/kvcache/offload_engine.py` - Add layerwise API
+- `nanovllm/engine/model_runner.py` - Pre-allocate buffers, refactor prefill/decode
+- `nanovllm/layers/attention.py` - Remove chunked prefill code
+- `nanovllm/utils/context.py` - Remove chunked prefill fields