phase 12: implement real continuous batching scheduler

Rewrote engine.rs from scratch: - Scheduler loop: admit → prefill → decode → finish → check new requests - Multiple sequences run concurrently (max_batch_size configurable) - Each sequence has independent GpuKVCache - Non-blocking try_recv() for new requests during decode iterations - Dynamic join: new requests enter batch immediately, don't wait for others Verified with concurrent test (tools/test_concurrent.py): - 3 concurrent requests: wall_time=3.8s, concurrency_ratio=2.82x ✓ - 5 concurrent requests: wall_time=6.1s, concurrency_ratio=4.04x ✓ - All outputs are coherent and correct Design doc (docs/12-continuous-batching.md) fully rewritten with: - Detailed scheduler loop pseudocode - Data structures (Sequence, Scheduler) - Acceptance criteria with specific test cases - Clear separation from Phase 13 (HTTP layer) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-22 13:44:26 +08:00
parent 7d05ececa0
commit d8493bd70f
5 changed files with 348 additions and 100 deletions
--- a/docs/12-continuous-batching.md
+++ b/docs/12-continuous-batching.md
@@ -2,100 +2,152 @@

 ## Goal

-实现 iteration-level 请求调度器，支持多请求并发执行和动态 batch 管理。这是 LLM serving 系统的核心调度逻辑。
+实现 iteration-level 请求调度，支持多个请求并发生成 token。核心能力：同时发 N 个请求，N 个请求同时产出 token，新请求可以在 mid-generation 加入 batch。

-## 核心概念
+## 为什么需要 Continuous Batching

-### Static Batching vs Continuous Batching
-
-**Static（朴素）**:
+**当前问题（串行）**：
 ```
-Batch 1: [req1, req2, req3] → 等所有完成才开始下一批
-问题: req1 10 token 就完了，req3 要 200 token → req1 的 slot 空转
+时间 →  [req1 prefill][req1 decode x 100][req2 prefill][req2 decode x 50]...
+GPU利用: ████████████████████████████████████████████████████████████████████
+         req2 等了 100 个 token 的时间才开始
 ```

-**Continuous（本阶段目标）**:
+**目标（continuous batching）**：
 ```
-Iteration 1: [req1, req2, req3]  → req1 完成! slot 释放
-Iteration 2: [req2, req3, req4]  → req4 立即填入
-每一个 iteration（一次 forward pass）重新决定哪些请求参与
+时间 →  [req1+req2 prefill][req1+req2 decode][req1 done, req3 加入][req2+req3 decode]...
+GPU利用: ████████████████████████████████████████████████████████████████████
+         req2 和 req1 同时推理，req3 在 req1 完成后立即加入
 ```

-## 核心组件
+## 核心设计

-### Sequence
+### 数据结构

 ```rust
 pub struct Sequence {
-    pub id: SeqId,
+    pub id: u64,
    pub prompt_tokens: Vec<u32>,
    pub generated_tokens: Vec<u32>,
-    pub status: SequenceStatus,
-    pub sampling_params: SamplingParams,
-    pub kv_cache_handle: KVCacheHandle,  // 该 seq 的 KV cache 资源
-    pub arrival_time: Instant,
-    pub output_sender: tokio::sync::mpsc::Sender<GenerateEvent>,
+    pub status: SeqStatus,
+    pub max_tokens: usize,
+    pub kv_cache: GpuKVCache,          // 每个 seq 独立的 KV cache
+    pub output_tx: mpsc::Sender<GenerateEvent>,
 }

-pub enum SequenceStatus {
-    Waiting,       // 等待调度
-    Prefilling,    // 正在 prefill
-    Decoding,      // 正在逐 token decode
-    Finished,      // 完成 (EOS / max_len)
+pub enum SeqStatus {
+    Waiting,     // 在队列中等待被 admit
+    Running,     // 正在参与 batch forward
+    Finished,    // EOS 或 max_tokens 达到
 }
-```

-### Scheduler
-
-```rust
 pub struct Scheduler {
-    waiting: VecDeque<Sequence>,      // 等待队列
-    running: Vec<Sequence>,           // 正在执行
-    max_batch_size: usize,            // 最大并发数
-    block_manager: BlockManager,      // KV cache 资源管理
+    waiting: VecDeque<Sequence>,
+    running: Vec<Sequence>,
+    max_batch_size: usize,             // 最大并发请求数
+    next_seq_id: u64,
 }
 ```

-### 调度循环
+### 调度循环（Engine 主循环）

 ```rust
 loop {
-    // 1. 回收已完成的 sequence，释放 KV cache
-    // 2. 从 waiting 中 admit 新请求（如果有空位+显存）
-    // 3. 对 running 中的所有 seq 做一步 forward
-    //    - 新加入的做 prefill
-    //    - 已在运行的做 decode
-    // 4. 对每个 seq 的 logits 做 sampling
-    // 5. 发送新 token / 完成信号
+    // Step 1: 回收已完成的 sequence
+    running.retain(|seq| seq.status != Finished);
+
+    // Step 2: Admit 新请求（如果 running < max_batch_size）
+    while running.len() < max_batch_size {
+        if let Some(seq) = waiting.pop_front() {
+            running.push(seq);
+        } else {
+            break;
+        }
+    }
+
+    if running.is_empty() {
+        // 没有任何工作，等待新请求
+        let new_req = request_rx.recv();  // blocking wait
+        waiting.push_back(new_req);
+        continue;
+    }
+
+    // Step 3: 分类 — 哪些需要 prefill，哪些需要 decode
+    let to_prefill: 新加入的 seq（generated_tokens 为空）
+    let to_decode: 已在运行的 seq
+
+    // Step 4: 执行
+    for seq in to_prefill {
+        // Prefill: 完整 prompt 一次 forward
+        model.forward_gpu_cache(&seq.prompt_tokens, &mut seq.kv_cache);
+        seq.status = Running;
+    }
+
+    // Decode: 每个 seq 独立做一步（当前不做 batch forward，留待优化）
+    for seq in to_decode {
+        let last_token = seq.last_generated_token();
+        let logits = model.forward_gpu_cache(&[last_token], &mut seq.kv_cache);
+        let next = sample_greedy(&logits);
+        seq.generated_tokens.push(next);
+        // 发送 token 给客户端
+        seq.output_tx.blocking_send(Token { id: next, text: decode(next) });
+        // 检查完成
+        if next == eos || seq.generated_tokens.len() >= seq.max_tokens {
+            seq.output_tx.blocking_send(Done);
+            seq.status = Finished;
+        }
+    }
+
+    // Step 5: 检查是否有新请求到达（non-blocking）
+    while let Ok(new_req) = request_rx.try_recv() {
+        waiting.push_back(new_req);
+    }
 }
 ```

-## 当前状态 (Phase 12 初版)
+### 关键设计决策

-当前实现是 **单请求顺序执行**（max_batch_size=1），是 continuous batching 的退化形式：
- 一次只处理一个请求
- 完成后才接受下一个
- 无 preemption、无 batching
+1. **每个 seq 独立 KV cache**：当前不做 batch forward（需要对齐 seq_len），而是每个 seq 独立调用 model.forward_gpu_cache。未来优化为 batched forward。

-这是合理的起步——先跑通单请求 E2E，后续扩展为真正的并发 batching。
+2. **Prefill 和 Decode 混合**：新加入的 seq 先 prefill（一次 forward），然后下一轮加入 decode batch。

-## 后续扩展 (Phase 15+)
+3. **Non-blocking request receive**：decode 循环中用 `try_recv()` 检查新请求，不阻塞推理。

-1. **多请求 batch forward**: 将多个 seq 的 token 拼接为一个 batch 输入
-2. **Prefill-Decode 分离**: prefill (compute-bound) 和 decode (memory-bound) 分开调度
-3. **Preemption**: 显存不足时暂停低优先级 seq
-4. **动态 batch size**: 根据 KV cache 使用量调整
+4. **max_batch_size**：受限于 GPU 显存（每个 seq 的 KV cache 占用）。Qwen3-8B 单卡 32GB，每个 seq 的 KV cache 约 256 tokens × 8 heads × 128 dim × 2(KV) × 2B = 1MB。可以并发 ~100 seq。实际受限于推理速度。

-## Test Plan
+## 与 Phase 13 (HTTP API) 的接口

- [x] 单请求 E2E: 提交请求 → 收到 token 流 → 完成信号
- [ ] (后续) 多请求并发: 提交多个请求，验证都能正确完成
- [ ] (后续) 短请求完成后新请求立即加入
+```
+HTTP Handler                              Engine Thread
+     │                                          │
+     │  ──── GenerateRequest ────────►          │
+     │       (prompt_tokens, max_tokens,        │
+     │        output_tx)                        │
+     │                                          │
+     │  ◄──── GenerateEvent (Token/Done) ────   │
+     │       (via tokio::sync::mpsc)            │
+     │                                          │
+```

-## Takeaways
+多个 HTTP handler 可以同时提交请求。Engine 线程内部通过 Scheduler 管理并发。

-1. **单请求是 continuous batching 的特殊情况 (batch_size=1)**：当前实现的 engine 循环已经是正确的调度结构——receive request → prefill → decode loop → done → next request。扩展为多请求只需在 decode loop 中处理多个 sequence。
+## 验收测试

-2. **Engine 在独立 OS thread 上跑是正确的设计**：GPU 操作是同步阻塞的（cudaDeviceSynchronize），如果放在 tokio runtime 中会 block 整个 async runtime。独立线程 + channel 通信是标准模式。
+必须通过以下测试才算 Phase 12 完成：

-3. **std::sync::mpsc::SyncSender(capacity=1) 实现了天然的背压**：当 engine 忙时，新请求会 block 在 channel send 上，不会积压。
+1. **并发 3 请求测试**：同时发 3 个请求，验证 3 个请求同时产出 token（不是串行等待）
+2. **吞吐量测试**：并发请求的总 token 吞吐量应接近单请求（因为单个 seq 的 decode 是串行的）
+3. **动态加入测试**：先发 1 个请求开始生成，过 2 秒再发第 2 个，验证第 2 个立即开始（不等第 1 个完成）
+4. **正确性测试**：并发请求的输出内容应与单独跑每个请求一致
+
+## 实现计划
+
+1. 重构 Engine：从 `while recv → generate` 改为 scheduler loop
+2. 每个 Sequence 持有独立的 GpuKVCache
+3. 调度循环实现 admit + prefill + decode + finish
+4. HTTP API 侧改为 unbounded channel（允许多请求同时提交）
+5. 编写并发测试脚本
+
+## 当前状态
+
+**未实现**。当前是 FIFO 串行，一次只处理一个请求。本文档是实现的设计规格。