从零构建 RWKV 批量推理服务器：2的幂次动态缩容、异步拷回与向量化采样

从零构建 RWKV 批量推理服务器：2的幂次动态缩容、异步拷回与向量化采样

引言

RWKV 是一种使用线性注意力机制的模型架构。与 Transformer 的 O(n²) 注意力不同，RWKV 推理时只需维护一个固定大小的隐藏状态（state），每个 token 的计算量恒定------推理复杂度 O(1)。

这个特性让它天生适合批量推理：Transformer 做 batch inference 需要精细管理不断膨胀的 KV Cache，而 RWKV 只需要把多个 state 摞在一起批处理。batch 越大，GPU 利用率越高，吞吐量近似线性增长。

但社区里一直没有开箱即用的批量推理服务。rwkv_lightning 有 state pool 的思路，Albatross 有 batch inference 的底层能力，Rapid-Sampling 支持向量化采样------但都得自己把这些零件组装起来，再包一层 HTTP 服务，再处理并发、任务排队、state 管理等等。

最近 GitHub 上做了一个项目，花了些时间把这件事做完了：RWKV-Server------一个兼容 OpenAI API 的线性注意力批量推理服务器。本文将走一遍这个项目的核心架构和几个值得展开的设计决策，希望能给同样在折腾 RWKV 部署的朋友一些参考。

---

目录

• [一、为什么批量推理是 RWKV 的强项？](#一、为什么批量推理是 RWKV 的强项？ "#%E4%B8%80%E4%B8%BA%E4%BB%80%E4%B9%88%E6%89%B9%E9%87%8F%E6%8E%A8%E7%90%86%E6%98%AF-rwkv-%E7%9A%84%E5%BC%BA%E9%A1%B9")
• 二、整体架构
• [三、调度器核心：Worker 槽位与声明式参数映射](#三、调度器核心：Worker 槽位与声明式参数映射 "#%E4%B8%89%E8%B0%83%E5%BA%A6%E5%99%A8%E6%A0%B8%E5%BF%83worker-%E6%A7%BD%E4%BD%8D%E4%B8%8E%E5%A3%B0%E6%98%8E%E5%BC%8F%E5%8F%82%E6%95%B0%E6%98%A0%E5%B0%84")
• 四、2的幂次动态缩容
• [五、脉冲式推理与 CUDA Stream 异步拷回](#五、脉冲式推理与 CUDA Stream 异步拷回 "#%E4%BA%94%E8%84%89%E5%86%B2%E5%BC%8F%E6%8E%A8%E7%90%86%E4%B8%8E-cuda-stream-%E5%BC%82%E6%AD%A5%E6%8B%B7%E5%9B%9E")
• [六、向量化 EOS 检测](#六、向量化 EOS 检测 "#%E5%85%AD%E5%90%91%E9%87%8F%E5%8C%96-eos-%E6%A3%80%E6%B5%8B")
• [七、Task 生命周期与上下文管理](#七、Task 生命周期与上下文管理 "#%E4%B8%83task-%E7%94%9F%E5%91%BD%E5%91%A8%E6%9C%9F%E4%B8%8E%E4%B8%8A%E4%B8%8B%E6%96%87%E7%AE%A1%E7%90%86")
• [八、提示词模板：max_tokens=0 的妙用](#八、提示词模板：max_tokens=0 的妙用 "#%E5%85%AB%E6%8F%90%E7%A4%BA%E8%AF%8D%E6%A8%A1%E6%9D%BFmax_tokens0-%E7%9A%84%E5%A6%99%E7%94%A8")
• [九、异步桥接：stream_callback 的 sync→async 转换](#九、异步桥接：stream_callback 的 sync→async 转换 "#%E4%B9%9D%E5%BC%82%E6%AD%A5%E6%A1%A5%E6%8E%A5stream_callback-%E7%9A%84-syncasync-%E8%BD%AC%E6%8D%A2")
• 十、日志与工程化细节
• 十一、实测性能
• 十二、配置与调优
• 十三、快速开始

---

一、为什么批量推理是 RWKV 的强项？

先看一张对比表：

	传统 LLM 推理服务	RWKV-Server
推理复杂度	O(n²) 注意力	O(1) 线性注意力
批量机制	Continuous Batching（分页管理 KV Cache）	状态复用 + 脉冲并行
用户感知	需手动管理 batch / queue	完全无感，提交即走
显存	KV Cache 随上下文线性增长	固定 state 大小，可预测
容量管理	静态 batch size	2的幂次动态缩容

核心区别在于 state 的管理方式。Transformer 的每个 token 需要跟所有历史 token 做注意力计算，所以 KV Cache 随上下文长度线性增长------做批量推理时你得精心控制每个序列的长度，否则显存就爆了。RWKV 的 state 是固定大小的（只跟模型维度有关，不随上下文增长），所以你拿到的是一张可预测的显存账单。

对于 RWKV7 2.9B 模型，每个任务约 11MB state 开销。256 个并发任务，state 总共约 2.5GB，加上模型权重约 6.5GB，一张 12GB 的 RTX 5070 Ti Laptop 就跑得动。

RWKV-Server 围绕三个原则构建：

1. 零感知 ------ 你不需要关心 batch 何时攒满、state 如何管理。提交 prompt，服务端自动完成一切。
2. 高吞吐 ------ 利用线性注意力特点，动态攒批 + 脉冲式生成，逼近 GPU 算力上限。
3. 弹性伸缩 ------ 根据负载自动调整 batch 容量（1→2→4→8→...→256），空闲时释放显存。

---

二、整体架构

五层架构，从上到下：

客户端 (OpenAI SDK / Chatbox / LobeChat / Open WebUI)
    │
    ▼
API 层 --- FastAPI :8000
    ├── /v1/chat/completions    (OpenAI 兼容)
    ├── /v1/completions         (文本补全)
    └── /v1/tasks/*              (私有任务管理，9个端点)
    │
    ▼
业务层 --- TaskManager
    ├── SQLite 持久化          (WAL 模式)
    ├── LRU CPU 缓存            (OrderedDict)
    └── 异步 I/O               (ThreadPoolExecutor + BoundedQueue)
    │
    ▼
调度层 --- DynamicScheduler / SimpleScheduler
    ├── 2的幂次动态缩容
    ├── Worker 槽位管理
    └── 声明式参数映射 (field_mappings)
    │
    ▼
引擎层 --- InferEngine + BatchSampler
    ├── 脉冲式批量 forward
    ├── CUDA Stream 异步拷回
    └── Per-sample 向量化采样
    │
    ▼
模型层 --- RWKV7 (Albatross 子模块) + Rapid-Sampling (向量化改造版)

关键模块一览：

模块	文件	核心职责
模型加载	loader.py	加载模型、tokenize、向量化 EOS 检测
调度器基类	scheduler_base.py	Worker 槽位、声明式 field_mappings、批量收集
动态调度器	scheduler_dynamic.py	2的幂次扩缩容、紧凑迁移
简化调度器	scheduler_simple.py	继承基类不加逻辑，供用户自定义
推理引擎	batch_engine.py	脉冲批量 forward + Stream 异步拷回
批量采样器	batch_sampler.py	向量化 per-sample 采样
任务管理	task/manager.py	单例、SQLite WAL、LRU 淘汰、异步写
OpenAI API	routes/v1/openai/	Chat Completion + Completion + /models
私有 API	routes/v1/rwkv/	9 个任务 CRUD 端点

---

三、调度器核心：Worker 槽位与声明式参数映射

调度器是整个系统的中枢。BaseScheduler 维护一个 Worker 字典，里面是预分配的 GPU 张量：

worker = {
    "last_tokens":     torch.zeros(B, dtype=torch.int32, device="cuda"),
    "max_tokens":      torch.zeros(B, dtype=torch.int32, device="cuda"),
    "tasks":           [None] * B,
    "generated_tokens": torch.zeros((B, buffer_size), dtype=torch.int32, device="cuda"),
    "shift_state":     ...,   # (num_layers, 2, B, dim)
    "wkv_state":       ...,   # (num_layers, B, head, 64, 64)
    "elapsed_t":       ...,   # (B,)
    "penalties":       ...,   # (B, vocab_size)
    "rand_state":      ...,   # (64*B,)
    "temperatures":    ...,   # (B,)
    "top_ps":          ...,   # (B,)
    "top_ks":          ...,   # (B,)
    "presence_penalties": ..., # (B,)
    "repetition_penalties": ..., # (B,)
    "penalty_decays":  ...,   # (B,)
    "stop_flags":      torch.zeros(B, dtype=torch.bool, device="cuda"),
    "mask":            torch.ones(B, dtype=torch.bool, device="cuda"),
}

mask=True 表示槽位空闲，mask=False 表示有任务在跑。

声明式 field_mappings 是 update_batch() 的精髓。不同于传统的逐字段硬编码赋值，这里用一张映射表声明每个字段如何从 Task 对象取出并注入到 Worker 张量中：

field_mappings = [
    ("stop_flags",        lambda s: s,                                   lambda t: None),     # 清零
    ("last_tokens",       lambda s: s,                                   lambda t: t.current_token),
    ("max_tokens",        lambda s: s,                                   lambda t: t.max_tokens),
    ("tasks",             lambda s: s,                                   lambda t: t),
    ("shift_state",       lambda s: (slice(None), slice(None), [s], slice(None)), lambda t: t.shift_state),
    ("wkv_state",         lambda s: (slice(None), [s], ...),             lambda t: t.wkv_state),
    ("elapsed_t",         lambda s: s,                                   lambda t: t.elapsed_t),
    ("penalties",         lambda s: s,                                   lambda t: t.penalties),
    ("rand_state",        lambda s: slice(s*64, (s+1)*64),               lambda t: t.rand_state),
    ("temperatures",      lambda s: s,                                   lambda t: t.temperature),
    ("top_ps",            lambda s: s,                                   lambda t: t.top_p),
    ("top_ks",            lambda s: s,                                   lambda t: t.top_k),
    # ... 更多字段
]

每个空闲槽位对应一个 Ready 状态的 Task。映射表依次遍历：取槽位索引 slot → 算出目标索引 → 从 Task 取 source 值 → 赋值到 Worker 张量中。通用、可读、易扩展------加一个新的 per-sample 参数只需在映射表里加一行。

---

四、2的幂次动态缩容

这是 DynamicScheduler 区别于 SimpleScheduler 的核心能力。流程如下：

flowchart TD A["_adjust_capacity()"] --> B["计算 total_needed = active + ready"] B --> C["target = next_power_of_two(total_needed)"] C --> D{target == current_capacity?} D -->|是| E[不操作] D -->|否| F["_compact_and_resize(target)"] F --> G[获取活跃槽位索引] G --> H["紧凑复制所有张量 (shift_state, wkv_state, penalties...)"] H --> I[复制 rand_state] I --> J[重建 tasks 列表] J --> K[重置 mask, 清零新增空闲槽位] K --> L["gc.collect() + torch.cuda.empty_cache()"] L --> M["完成: capacity N → M"]

DynamicScheduler.run() 的关键流程：

def run(self):
    self._adjust_capacity()  # 首次扩容（能装下所有 ready 任务）
    while self.tasks:
        self.update_batch()          # 注入任务
        self._adjust_capacity()      # 中间调整（如有新增任务）
        self.update_batch()          # 再次注入（空间更多了）
        # 实际推理
        self.engine.generate(self.worker, self.mask, self.worker["stop_flags"])
        self._collect()              # 收集结果 + 清理已完成任务

注意 _adjust_capacity 被调用了两次------一次在 run 循环外（确保初始容量能装下所有待处理任务），一次在 update_batch 之间（如果有新任务加入且容量不够，先扩容再继续注入）。这种设计避免了"任务排队等扩容"的尴尬。

为什么是 2 的幂而不是任意值？

1. CUDA 内存对齐 ：shift_state 形状为 (num_layers, 2, B, dim)，其中 B 为 batch 维度。当 B 是 2 的幂次时，CUDA 的访存对齐最优
2. 减少碎片化 ：从 1→2→4→...→256 只需 9 个级别。如果逐量调整（1→2→3→...→256），频繁的 cudaMalloc/cudaFree 会产生大量显存碎片
3. 紧凑迁移成本可控：幂次跳跃意味着扩容频率低，每次迁移的分摊成本小
4. 状态向量效率 ：wkv_state 形状 (num_layers, B, head, 64, 64)，B 为幂次时内层 64×64 矩阵的批量 GEMM 效率最高

紧凑迁移的具体操作：把所有活跃槽位的张量统一往低位复制。_compact_and_resize 中维护了一张 tensor_moves 表，跟 field_mappings 思路一致但操作不同------这里是低位紧凑复制而非注入新值。如果新容量 > 旧容量，在尾部分配更多零初始化的槽位；如果变小，直接裁剪。

---

五、脉冲式推理与 CUDA Stream 异步拷回

脉冲（Pulse）机制

每一次 engine.generate() 调用生成 buffer_size 个 token（默认 32）。这就是一个 pulse。每个 pulse 内部循环：

for step in range(buffer_size):
    logits = self.model.forward(last.unsqueeze(-1), [s0, s1, s3])
    next_tokens = self.sampler.sample(logits, ...)  # (B,) 向量化采样
    gen[:, step] = next_tokens                        # 记录生成 token
    new_finish = self.eos_fn(last, next_tokens)       # (B,) 向量化 EOS
    last.copy_(next_tokens)                           # 更新 last_tokens
    mask |= new_finish                                # 标记已完成槽位

buffer_size 的选择是个权衡------太小（如 8）意味着频繁的 Python↔CUDA 调度开销；太大（如 128）意味着高并发时 GPU 压力更大。实测 16-64 比较合理，默认 32。

CUDA Stream 异步拷回

任务完成时需要把 state 拷回 CPU。但如果走默认 stream，这个 D2H 拷贝会阻塞后续推理------所有 CUDA 操作默认在同一流上串行执行。

解决：预分配一批 pin_memory（页锁定内存，D2H 拷贝最快）作为中转缓冲区，拷贝走独立的 copy_stream：

class InferEngine:
    def __init__(self, ...):
        s, w, e = self.model.zero_state(1)
        # 预分配 pin_memory 缓冲区（每个任务一个，复用）
        self.temp_shift_state = s.cpu().pin_memory()
        self.temp_wkv_state = w.cpu().pin_memory()
        self.temp_elapsed_t = e.cpu().pin_memory()
        self.temp_penalties = torch.zeros(vocab_size).pin_memory()
        self.temp_rand_state = torch.zeros(64, dtype=torch.int8).pin_memory()
        self.copy_stream = torch.cuda.Stream()

    def generate(self, ...):
        for step in range(self.buffer_size):
            # ... forward + sample + EOS 检测 ...

            # 完成的槽位：异步拷回 state
            for idx in finish_indices:
                task = tasks[idx]
                with torch.cuda.stream(self.copy_stream):
                    task.shift_state.copy_(self.temp_shift_state)    # GPU→CPU (async)
                    task.wkv_state.copy_(self.temp_wkv_state)
                    task.elapsed_t.copy_(self.temp_elapsed_t)
                    task.penalties.copy_(self.temp_penalties)
                    task.rand_state.copy_(self.temp_rand_state)

            # 原地掩码：finish 的槽位 state 清零（不影响其他槽位）
            shift_state[:, :, :cur_batch, :] *= ~new_finish.view(1, 1, cur_batch, 1)
            wkv_state[:, :cur_batch, ...]   *= ~new_finish.view(1, cur_batch, 1, 1, 1)

# 注意：不需要 torch.cuda.synchronize(self.copy_stream) ------ 推理主流程不等拷贝完成

上一个 pulse 的 state 还在异步拷贝，下一个 pulse 已经开始了。而且 torch.cuda.current_stream() 上的计算不依赖 copy_stream 上的操作，所以是真正的零阻塞。

---

六、向量化 EOS 检测

RWKV 的结束符设计不同于 Transformer 的单个 <eos> token。它的对话格式是 User:...\n\nAssistant:...<|endoftext|>，实际触发结束的条件有三种：

模式	条件	含义
单 token 匹配	cur_token == 0	标准结束符 `<
\n\n + User	prev == 261 && cur == 24281	模型认为下一轮用户要发言了
\n + User	prev == 11 && cur == 24281	同上但只有单个换行

全部在 GPU 上向量化完成：

def batch_is_eos(self, prev_tokens, cur_tokens):
    """prev_tokens: (B,) int32, cur_tokens: (B,) int32"""
    cond1 = cur_tokens == self._eos_single_token          # token 0
    cond2 = (prev_tokens == 261) & (cur_tokens == 24281)  # \n\n → User
    cond3 = (prev_tokens == 11)  & (cur_tokens == 24281)  # \n  → User
    return cond1 | cond2 | cond3  # (B,) bool

没有 Python 循环，没有 CPU 回传------三个条件都是 GPU 上的逐元素布尔运算，延迟几乎为零。

---

七、Task 生命周期与上下文管理

每个 Task 对象在四种状态之间流转：

  new_task()
      │
      ▼
  [PREFILL] ──prefill完成──▶ [READY] ──被调度器pick──▶ [RUNNING]
                                                            │
                                              EOS / max_tokens / stop()
                                                            │
                                                            ▼
                               [FINISHED] ◀── continue_gen() 可回到 READY

每个状态的含义：

• PREFILL：正在 prefill prompt（CPU 上完成，非 GPU）
• READY：prefill 完成，state 在 CPU 上，等待调度器拾取
• RUNNING：正在 GPU 上推理，state 在 GPU 上
• FINISHED：生成完成，state 已拷回 CPU，结果已收集

关键设计在于 Python 上下文管理器 ( with task: ) 自动处理 GPU↔CPU 迁移：

class Task:
    def __enter__(self):
        self.prepare()    # 标记 RUNNING，记录 run_time，将 state 移到 GPU
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cpu()        # 将 state 移回 CPU

# 在 update_batch 中的使用：
for slot, task in zip(free_slots, ready_tasks):
    with task:  # 自动 GPU 迁移
        # 把 task 的 state 拷贝到 worker 的 GPU 张量中
        worker["shift_state"][slot] = task.shift_state
        worker["wkv_state"][slot] = task.wkv_state
        # ...
    # __exit__ 自动触发 cpu()，state 移回 CPU

这个模式保证了 state 在 GPU 上的时间窗口精确可控------只在拷贝那几行代码期间占用 GPU 内存，用完立即回 CPU。避免了"忘记手动 cpu()"导致的显存泄漏。

---

八、提示词模板：max_tokens=0 的妙用

这是整个项目中最实用的设计之一------一个简单的参数设定，打开了提示词工程的大门。

场景：你有一个非常好用的系统提示词（如"你是一个专业的 Python 后端工程师，请遵循以下编码规范..."），每次请求都要带上它。常规做法是每次 prefill------重复计算同一个前缀。

RWKV-Server 的做法：

# Step 1：只 prefill，不生成 → state 入库
template_task = scheduler.new_task(
    prompt="User: 你是一个资深 Python 工程师...\n\nAssistant: 好的，我开始写。\n\n",
    max_tokens=0,    # ★ 核心：只跑 prefill
)
scheduler.run()                         # prefill 完成
task_manager.put_task("_coder", template_task)  # 持久化到 SQLite

# Step 2：每次使用时 fork → 追加真正的任务 prompt
new_id = task_manager.fork_template("_coder")   # state 自动继承
task = task_manager.get_task(new_id)
task.continue_gen(
    "User: 写一个线程安全的单例模式实现。\n\nAssistant:",
    max_tokens=500,
)
scheduler.add_task(task)               # 从 prefill 后的 state 开始推理

模板 ID 以 _ 开头，被视为只读------不受容量限制淘汰，不会被覆盖。模板存储在 SQLite（WAL 模式），重启服务依然可用。

进阶玩法：预设多个角色，agent 在某个推理分支 fork 不同模板，甚至并行：

roles = {
    "_coder":       "...",
    "_translator":  "...",
    "_writer":      "...",
}
for name, prompt in roles.items():
    t = scheduler.new_task(prompt=prompt, max_tokens=0)
    scheduler.run()
    task_manager.put_task(name, t)

# 运行时 fork，state 复用，零 prefill 开销
coding_task  = task_manager.fork_template("_coder")
writing_task = task_manager.fork_template("_writer")

---

九、异步桥接：stream_callback 的 sync→async 转换

调度器和推理引擎运行在同步线程中（threading.Thread），但 SSE 流式输出需要在 FastAPI 的 async 上下文中工作。这里的桥接方案简洁而关键：

def stream_callback():
    loop = asyncio.get_event_loop()
    queue: asyncio.Queue = asyncio.Queue()
    _finished = False

    def callback(data):
        if _finished:
            return
        loop.call_soon_threadsafe(queue.put_nowait, data)  # 同步→异步

    async def generator():
        while True:
            chunk = await queue.get()
            if chunk is None:
                break
            yield chunk

    def finish(data=None):
        nonlocal _finished
        if _finished:
            return
        _finished = True
        loop.call_soon_threadsafe(queue.put_nowait, None)
        loop.call_soon_threadsafe(queue.put_nowait, None)
        loop.call_soon_threadsafe(queue.put_nowait, None)  # 3次确保

    return generator(), callback, finish

三个返回值：

• generator()：异步生成器，供 SSE 端点消费
• callback(data)：同步线程调用，把 token 块推入 Queue
• finish(data)：同步线程调用，推送结束信号（3 个 None，防止 SSE 提前关闭）

call_soon_threadsafe 是 Python asyncio 提供的官方桥接 API------在任意线程中安全地将回调调度到事件循环中执行。比手动 threading.Queue + 轮询优雅得多，且完全线程安全。

---

十、日志与工程化细节

Loguru 劫持 uvicorn/starlette 日志

FastAPI 底层用 uvicorn 和 starlette，它们使用 Python 标准库的 logging 模块。如果不处理，终端会一半 loguru 格式、一半标准库格式，很混乱。

通过 InterceptHandler（继承 logging.Handler）劫持标准库日志，重定向到 loguru：

class InterceptHandler(logging.Handler):
    def emit(self, record):
        # 获取调用栈深度，确保 loguru 正确追踪源文件
        frame, depth = inspect.currentframe(), 0
        while frame and (depth == 0 or frame.f_code.co_filename == logging.__file__):
            frame = frame.f_back
            depth += 1
        logger.opt(depth=depth, exception=record.exc_info).log(level, record.getMessage())

统一格式：HH:mm:ss.SS | LEVEL | module | message，所有组件输出一致。

其他细节

• 调度器日志模块名 ：scheduler_base.py 动态绑定 module_name，根据类名自动区分 scheduler.dynamic / scheduler.simple，部署调试时一眼看出是哪个调度器在跑
• TaskManager 单例模式：双重检查锁定（double-checked locking），确保全局只有一个实例，SQLite 连接安全
• BoundedThreadPoolExecutor：自定义线程池包装，异步写 DB 时有界队列防止无限堆积
• Locust 压测脚本：开箱即用，自动记录 TTFT（首包延迟）、SSE 流式总耗时、平均轮询次数等指标

---

十一、实测性能

RTX 5070 Ti Laptop（12GB 显存），RWKV7 2.9B 模型。256 个长 prompt 任务并发，max_tokens=2000，temperature=1.0，buffer_size=32：

 RWKV 推 理 压 力 测 试
 模型: 2.9B · 显卡: RTX 5070 Ti Laptop GPU · 任务数: 256 · 显存: 9.0 GB
────────────────────────────────────────────────
 输入 tokens:  20,140       输出 tokens:  545,440
 总吞吐量:     1,973 tok/s  输出吞吐:     1,903 tok/s
 平均每任务输出: 2,147 tokens
────────────────────────────────────────────────
 峰值: Iter 2 cap=256 active=256 speed=5168.75 tok/s per_task=20.19 tok/s

指标	数值
峰值吞吐	5,169 tok/s（256 并发满载）
全程平均吞吐	1,973 tok/s
单任务 256 并发	~20 tok/s
单任务低并发	~55-80 tok/s
显存占用	9.0 GB（模型 6.5GB + 256 state ~2.5GB）
总耗时	286s（256 任务 × 平均 ~2,147 token 输出）

每个任务约 11MB state 开销。7.2B 模型（约 14.4GB 权重）在 24GB 显存上跑 256 并发理论可行。

---

十二、配置与调优

所有配置通过 RWKV_ 前缀环境变量或直接修改 server/config.py（Pydantic Settings）：

环境变量	默认值	说明
RWKV_MODEL_PATH	server/model/rwkv7-g1g-2.9b-xxx.pth	模型路径
RWKV_MAX_BATCH_SIZE	256	最大并发数
RWKV_BUFFER_SIZE	32	每个 pulse 生成 token 数
RWKV_DEFAULT_MAX_TOKENS	4096	默认最大生成 token
RWKV_DEFAULT_TEMPERATURE	1.0	采样温度
RWKV_DEFAULT_TOP_P	0.3	Top-P 核采样
RWKV_DEFAULT_TOP_K	50	Top-K 采样
RWKV_DEFAULT_PRESENCE_PENALTY	2.0	存在惩罚（鼓励多样性）
RWKV_DEFAULT_REPETITION_PENALTY	0	重复惩罚
RWKV_DEFAULT_PENALTY_DECAY	1.0	惩罚衰减
RWKV_VERBOSE	true	日志级别
RWKV_TASK_DB_PATH	rwkv_tasks.db	SQLite 路径
RWKV_TASK_DEFAULT_CPU_CAPACITY	300	LRU 缓存容量
RWKV_TASK_DB_MAX_SIZE	10000	SQLite 最大非模板任务

调优建议：

• BUFFER_SIZE：越大单次 pulse 生成越多，但高负载下 GPU 压力更大。建议 16-64
• MAX_BATCH_SIZE：2.9B 模型 256，7.2B 模型 256（每个任务 ~11MB）
• PRESENCE_PENALTY：创意写作建议 2.0+，代码生成建议 0.5-1.0
• TOP_P：默认 0.3（偏保守），创意场景可调高到 0.7-0.9

---

十三、快速开始

# 克隆（含子模块）
git clone --recurse-submodules https://github.com/AUXStar/RWKV-Server.git
cd RWKV-Server

# 安装依赖（需要 uv）
uv sync

# 设置模型路径（或直接改 server/config.py）
$env:RWKV_MODEL_PATH = "path/to/rwkv7-xxx.pth"   # Windows
export RWKV_MODEL_PATH="path/to/rwkv7-xxx.pth"    # Linux/macOS

# 启动
python run.py

启动后：

12:01:23.45 | INFO    | scheduler.dynamic  | Init capacity = 1 (max=256)
12:01:23.56 | INFO    | api                | Uvicorn running on http://0.0.0.0:8000

Python 客户端一行代码：

from openai import OpenAI
client = OpenAI(base_url="http://localhost:8000/v1", api_key="sk-local")
response = client.chat.completions.create(
    model="rwkv-7",
    messages=[{"role": "user", "content": "你好！"}],
    temperature=0.7, max_tokens=500,
)
print(response.choices[0].message.content)

---

结语

RWKV-Server 还在早期阶段，但核心骨架已经完整------调度器、引擎、采样器、任务管理、API 层都是解耦的。SimpleScheduler 只继承了基类没有加任何逻辑，就是留给想自定义调度策略的开发者用的。

欢迎试用和反馈。如果你在 4090 或更高配置上跑出了数据，期待你的跑分贡献 🚀

GitHub：github.com/AUXStar/RWK...