RWKV-API

RWKV-Server 的官方 Python SDK，以最优雅的方式调用 RWKV 推理。

无需关心批量调度、状态管理、GPU 显存分配——只需创建一个 Task，RWKV-Server 会在后台自动完成一切。

配套服务端

本 SDK 对应 RWKV-Server 的 Task API。服务端基于 RWKV7 线性注意力架构，支持动态批量推理、2的幂次自动扩缩容、模板持久化等高级特性，单机可达 11,081 tok/s 峰值吞吐。

安装

pip install rwkv-api

快速开始

同步调用

from rwkv_api import Client

client = Client("http://localhost:8000")

# 创建任务并等待结果 —— 服务端自动批量调度
task = client.create("Hello world", max_tokens=50)
result = task.wait()
print(result.result)

# FIM 代码补全
result = client.fim(
    prefix="def add(a, b):\n    ",
    suffix="\n    return result",
    max_tokens=30
)
print(result.result)

异步调用

import asyncio
from rwkv_api import AsyncClient

async def main():
    client = AsyncClient("http://localhost:8000")
    task = await client.create("Hello world", max_tokens=50)
    result = await task.wait()
    print(result.result)

asyncio.run(main())

实时流式（创建时）

from rwkv_api import Client

client = Client("http://localhost:8000")

# 方式1：create_stream 显式流式
for chunk in client.create_stream("Tell me a story", max_tokens=200):
    print(chunk, end="", flush=True)

# 方式2：create(stream=True) 自动分发
for chunk in client.create("Tell me a story", max_tokens=200, stream=True):
    print(chunk, end="", flush=True)

对已创建任务订阅流式

from rwkv_api import Client

client = Client("http://localhost:8000")

# 先创建任务（非流式）
task = client.create("Hello", max_tokens=200)

# 稍后订阅实时流式输出
for chunk in task.stream():
    print(chunk, end="", flush=True)

# 或者直接用 task_id
for chunk in client.stream_task(task.task_id):
    print(chunk, end="", flush=True)

异步实时流式

import asyncio
from rwkv_api import AsyncClient

async def main():
    client = AsyncClient("http://localhost:8000")

    # 创建时流式
    async for chunk in client.create_stream("Tell me a story", max_tokens=200):
        print(chunk, end="", flush=True)

    # 对已创建任务订阅流式
    task = await client.create("Hello", max_tokens=200)
    async for chunk in task.stream():
        print(chunk, end="", flush=True)

asyncio.run(main())

Task 对象 —— 优雅调用 RWKV 推理的核心

SDK 的核心是 Task 对象。你只需关注要生成什么，所有底层细节（批量调度、状态生命周期、GPU 显存管理）都由服务端自动处理：

task = client.create("Hello", max_tokens=50)

# 属性
print(task.task_id)       # "TASK_xxx"
print(task.finished)      # False
print(task.result)        # 生成结果（完成后可用）
print(task.speed)         # tok/s
print(task.prefill_time)  # prefill 耗时（秒）
print(task.gen_time)      # 生成耗时（秒）

# 操作
result = task.wait()           # 阻塞等待完成，返回 self
task.stop()                    # 停止生成
forked = task.fork()           # Fork 新任务
continued = task.continue_()   # 继续生成
tpl = task.as_template()       # 转为模板
task.delete()                  # 删除任务

# 流式订阅（任务创建后仍可实时获取输出）
for chunk in task.stream():
    print(chunk, end="", flush=True)

流式说明

SDK 提供两种流式方式：

注意：

• create(stream=True) 返回生成器，不是 Task 对象
• task.stream(pos=N) 从第 N 个 token 开始（重连时传已读数量避免重复）
• 如果任务已完成，stream() 一次性返回完整结果

异常处理

from rwkv_api import Client
from rwkv_api.exceptions import (
    RWKVValidationError,   # 422 参数错误
    TaskNotFoundError,      # 404 任务不存在
    RWKVServerError,        # 500 服务端错误
    TimeoutError,           # 等待超时
    ConnectionError,        # 连接失败
)

client = Client("http://localhost:8000")

try:
    task = client.create("Hello", max_tokens=50)
    result = task.wait(timeout=10)
except RWKVValidationError as e:
    print(f"参数错误: {e}")
except TimeoutError:
    print("等待超时")
except RWKVServerError as e:
    print(f"服务端错误: {e}")

API 参考

Client / AsyncClient

Task / AsyncTask

生成参数

开发

git clone https://github.com/AUXStar/RWKV-API.git
cd RWKV-API
pip install -e ".[dev]"
pytest tests/ -v

许可证

MIT