RWKV pip 使用指南

import os, time
import numpy as np
import torch

# 环境变量设置：必须在导入 rwkv 库之前设置
os.environ["RWKV_V7_ON"] = '1'       # 显式开启 RWKV-v7 模型支持（如果你用的是 v6 模型，请注释掉或设为 0）
os.environ['RWKV_JIT_ON'] = '1'      # 开启 JIT (Just-In-Time) 编译，加速算子加载
os.environ["RWKV_CUDA_ON"] = '1'     # 开启 CUDA 自定义算子（需系统安装了 CUDA 编译器，速度显著提升）

from rwkv.model import RWKV
from rwkv.utils import PIPELINE

# 初始化模型
# strategy='cuda fp16': 使用 GPU (cuda) 进行计算，精度为 fp16
model = RWKV(model='/mnt/e/RWKV-Runner/models/rwkv7-g1a-0.1b-20250728-ctx4096', strategy='cuda fp16')

# 初始化 Pipeline
# 用于将文本转换为 token id (encode) 以及将 token id 转换回文本 (decode)
pipeline = PIPELINE(model, "rwkv_vocab_v20230424")

# 生成参数设置
LENGTH_PER_TRIAL = 256  # 每次生成的 token 长度
TEMPERATURE = 1.0       # 温度：越高越随机，越低越确定
TOP_P = 0.0             # Nucleus 采样：0.0 通常意味着贪婪采样（或由库的具体实现决定，RWKV中 0 通常指argmax）

# 提示词 (Prompt)
prompt = "User: simulate SpaceX mars landing using python\n\nAssistant: <think"

print('='*80 + '\nSlow inference\n' + '='*80)
print(prompt, end="")

all_tokens = []
out_last = 0

# 1. 预处理 Prompt (Prefill)
# 将提示词编码并输入模型，获取初始的 logits (out) 和状态 (state)
# state 包含了对 prompt 的所有记忆
out, state = model.forward(pipeline.encode(prompt), None)

times = []
all_times = []
t000 = time.perf_counter()

# 2. Token 生成循环
for i in range(LENGTH_PER_TRIAL):
    t00 = time.perf_counter()
    
    # 采样：根据模型输出的 logits 选择下一个 token
    token = pipeline.sample_logits(out, temperature=TEMPERATURE, top_p=TOP_P)
    all_tokens += [token]

    # 解码并打印：将生成的 token 实时转回文本显示
    tmp = pipeline.decode(all_tokens[out_last:])
    if '\ufffd' not in tmp: # 过滤掉无效的 UTF-8 字符（通常是多字节字符未接收完整时出现）
        print(tmp, end="", flush=True) 
        out_last = i+1    

    torch.cuda.synchronize() # 等待 GPU 完成上述操作，确保计时准确
    t0 = time.perf_counter()

    # 模型前向传播 (Forward)
    # 将选出的 token 和上一步的 state 输入模型，计算下一个 token 的概率
    out, state = model.forward(token, state)

    torch.cuda.synchronize()
    t1 = time.perf_counter()
    
    # 记录时间
    times.append(t1 - t0)       # 纯模型计算时间
    all_times.append(t1 - t00)  # 包含采样和 Python 逻辑的总时间

# 计算并打印速度统计
times = np.percentile(times, 50)      # 取中位数
all_times = np.percentile(all_times, 50)
print(f'\n\nToken/s = {round(1/times,2)} (forward), {round(1/all_times,2)} (full)')

print('='*80 + '\nFast inference (CUDAGraph, requires rwkv pip pkg v0.8.31+)\n' + '='*80)
print(prompt, end="")

all_tokens = []
out_last = 0

# 重置状态，获取一个全零的初始状态结构
state = model.generate_zero_state()

#  关键步骤：分配静态显存 (Static Memory) 
# CUDA Graph 要求输入和输出的内存地址必须固定，不能变动。

# 1. 静态输入 Tensor (存放当前 Token 的 Embedding)
static_input = torch.empty((model.n_embd), device="cuda", dtype=torch.half)

# 2. 静态状态 Tensor (存放 RNN 隐状态)
# 需要两组：in 用于输入上一刻状态，out 用于输出当前刻状态
static_state_in = [torch.empty_like(x, device="cuda") for x in state]
static_state_out = [torch.empty_like(x, device="cuda") for x in state]

# 3. 静态输出 Tensor (存放预测下一个词的 Logits)
static_output = torch.empty((model.args.vocab_size), device="cuda", dtype=torch.half)

#  关键步骤：录制图 (Graph Capture) 
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    # 在这个上下文管理器中，执行一次特殊的 forward 操作。
    # GPU 不会真的计算数据，而是"记住"了所有的计算步骤和依赖关系。
    # model.forward_one_alt 是专门为 CUDAGraph 优化的单步前向函数
    static_output, static_state_out = model.forward_one_alt(static_input, static_state_in)

# 1. 预处理 Prompt (Prefill)
# Prompt 处理阶段无法使用 CUDAGraph（因为长度不固定），所以使用常规 forward
out, state = model.forward(pipeline.encode(prompt), state)

# 2. 将 Prompt 计算后的状态复制到静态内存中，为 Graph 运行做准备
for i in range(len(state)):
    static_state_in[i].copy_(state[i])
    
# 将初始输出也复制进去（虽然如果是 greedy decoding 可能只用 state 就够，但保持一致性）
static_output.copy_(out)

times = []
all_times = []
t000 = time.perf_counter()

# 3. 快速生成循环
for i in range(LENGTH_PER_TRIAL):
    t00 = time.perf_counter()
    
    # 采样 (Sample) - 注意：这里直接从 static_output 读取数据
    token = pipeline.sample_logits(static_output, temperature=TEMPERATURE, top_p=TOP_P)
    all_tokens += [token]

    # 解码打印 (Decode)
    tmp = pipeline.decode(all_tokens[out_last:])
    if '\ufffd' not in tmp:
        print(tmp, end="", flush=True)
        out_last = i+1

    torch.cuda.synchronize()
    t0 = time.perf_counter()

    #  核心加速步骤 
    
    # A. 填充输入：直接从 Embedding 表中取出对应的向量，填入 static_input
    # 注意：这里我们不传 token id，而是直接传向量，减少图内部的查表开销
    static_input.copy_(model.z['emb.weight'][token])
    
    # B. 回放图：一键执行整个模型的前向计算
    g.replay()
    
    # C. 更新状态：将输出状态 (out) 复制回 输入状态 (in)，形成 RNN 的循环
    # 这样下一次 replay 时，使用的就是更新后的状态了
    for n in range(len(state)):
        static_state_in[n].copy_(static_state_out[n])
    
    # --
    
    torch.cuda.synchronize()
    t1 = time.perf_counter()
    times.append(t1 - t0)
    all_times.append(t1 - t00)

# 打印统计
times = np.percentile(times, 50)
all_times = np.percentile(all_times, 50)
print(f'\n\nToken/s = {round(1/times,2)} (forward), {round(1/all_times,2)} (full) (note: very inefficient sample_logits)')

########################################################################################################
# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
########################################################################################################

print("RWKV Chat Simple Demo")  # 打印一个简单的消息，表明这是 RWKV 聊天的简单演示。
import os, copy, types, gc, sys, re  # 导入操作系统、对象复制、类型、垃圾回收、系统、正则表达式等包
import numpy as np  # 导入 numpy 库
from prompt_toolkit import prompt  # 从 prompt_toolkit 导入 prompt，用于命令行输入
import torch  # 导入 pytorch 库

# 优化 PyTorch 设置，允许使用 tf32 
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_tf32 = True

# os.environ["RWKV_V7_ON"] = '1' # 启用 RWKV-7 模型
os.environ["RWKV_JIT_ON"] = "1" # 启用 JIT 编译
os.environ["RWKV_CUDA_ON"] = "0"  # 禁用原生 CUDA 算子，改成 '1' 表示启用 CUDA 算子（速度更快，但需要 c++ 编译器和 CUDA 库）

from rwkv.model import RWKV  # 从 RWKV 模型库中导入 RWKV 类，用于加载和操作 RWKV 模型。
from rwkv.utils import PIPELINE  # 从 RWKV 工具库中导入 PIPELINE，用于数据的编码和解码

args = types.SimpleNamespace()

args.strategy = "cuda fp16"  # 模型推理的设备和精度，使用 CUDA （GPU）并采用 FP16 精度
args.MODEL_NAME = "E://RWKV-Runner//models//rwkv-final-v6-2.1-1b6"  # 指定 RWKV 模型的路径，建议写绝对路径

# STATE_NAME = None # 不使用 State

# 指定要加载的 State 文件路径。
STATE_NAME = "E://RWKV-Runner//models//rwkv-x060-eng_single_round_qa-1B6-20240516-ctx2048"  # 指定要加载的自定义 State 文件路径。

# 设置模型的解码参数
GEN_TEMP = 1.0
GEN_TOP_P = 0.3
GEN_alpha_presence = 0.5
GEN_alpha_frequency = 0.5
GEN_penalty_decay = 0.996

# 判断是否加载了一个 State 文件。如果指定了 State ，则调整生成参数，使回答的效果更好。
if STATE_NAME != None:
    GEN_TOP_P = 0.2
    GEN_alpha_presence = 0.3
    GEN_alpha_frequency = 0.3

CHUNK_LEN = 256  # 对输入进行分块处理

print(f"Loading model - {args.MODEL_NAME}")# 打印模型的加载消息
model = RWKV(model=args.MODEL_NAME, strategy=args.strategy)  # 加载 RWKV 模型。
pipeline = PIPELINE(model, "rwkv_vocab_v20230424")  # 初始化 PIPELINE ，使用 RWKV-World 词表处理输入和输出的编码/解码。

model_tokens = []
model_state = None

# 如果指定了 STATE_NAME，则加载自定义 State 文件，并初始化模型 State
if STATE_NAME != None:
    args = model.args  # 获取模型参数
    state_raw = torch.load(STATE_NAME + '.pth')  # 从指定的 State 文件中加载 State 数据
    state_init = [None for i in range(args.n_layer * 3)]  # 初始化状态列表
    for i in range(args.n_layer): #开始循环，遍历每一层。
        dd = model.strategy[i]  # 获取模型每一层的加载策略
        dev = dd.device  # 获取每一层的加载设备（如 GPU）
        atype = dd.atype  # 获取每一层的数据类型（FP32/FP16 或 int8 等）
        # 初始化模型的状态
        state_init[i*3+0] = torch.zeros(args.n_embd, dtype=atype, requires_grad=False, device=dev).contiguous()
        state_init[i*3+1] = state_raw[f'blocks.{i}.att.time_state'].transpose(1,2).to(dtype=torch.float, device=dev).requires_grad_(False).contiguous()
        state_init[i*3+2] = torch.zeros(args.n_embd, dtype=atype, requires_grad=False, device=dev).contiguous()
    model_state = copy.deepcopy(state_init)  # 复制初始化的状态

def run_rnn(ctx):
    # 定义两个全局变量，用于更新 token 和模型状态（state）
    global model_tokens, model_state
    ctx = ctx.replace("\r\n", "\n")  # 将文本中的 CRLF（Windows 系统的换行符）转换为 LF（Linux 系统的换行符）
    tokens = pipeline.encode(ctx)  # 基于 RWKV 模型的词汇表，将文本编码为 tokens
    tokens = [int(x) for x in tokens]  # 将 tokens 转换为整数（int）列表，确保类型一致性
    model_tokens += tokens  # 将 tokens 添加到全局的模型 token 列表中

    while len(tokens) > 0:  # 使用一个 while 循环执行模型前向传播，直到所有 tokens 处理完毕
        out, model_state = model.forward(tokens[:CHUNK_LEN], model_state)  # 模型前向传播，处理大小为 CHUNK_LEN 的 token 列表，并更新模型状态
        tokens = tokens[CHUNK_LEN:]  # 移除已处理的 tokens 块，并继续处理剩余的 tokens

    return out  # 返回模型的 prefill 结果

# 如果没有加载自定义 State ，则使用初始提示进行对话
if STATE_NAME == None:
    init_ctx = "User: hi" + "\n\n"
    init_ctx += "Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it." + "\n\n"
    run_rnn(init_ctx)  # 运行 RNN 模式对初始提示文本进行 prefill
    print(init_ctx, end="")  # 打印初始化对话文本

# 从用户输入中读取消息、循环生成下一个 token 
while True:
    msg = prompt("User: ")  # 从用户输入中读取消息，存到 msg 变量
    msg = msg.strip()  # 使用 strip 方法去除消息的首尾空格
    msg = re.sub(r"\n+", "\n", msg)  # 替换多个换行符为单个换行符
    if len(msg) > 0:  # 如果处理完后，用户输入的消息非空
        occurrence = {}  # 使用 occurrence 字典这个字典用于记录每个 token 在生成上下文中出现的次数，等会用在实现重复惩罚（Penalty）
        out_tokens = []  # 使用 out_tokens 列表记录即将输出的 tokens
        out_last = 0  # 用于记录上一次生成的 token 位置

        out = run_rnn("User: " + msg + "\n\nAssistant:")  # 将用户输入拼接成 RWKV 数据集的对话格式，进行 prefill  
        print("\nAssistant:", end="")  # 打印 "Assistant:" 标签

        for i in range(99999):  
            for n in occurrence: 
                out[n] -= GEN_alpha_presence + occurrence[n] * GEN_alpha_frequency  # 应用存在惩罚和频率惩罚参数
            out[0] -= 1e10  # 禁用 END_OF_TEXT 

            token = pipeline.sample_logits(out, temperature=GEN_TEMP, top_p=GEN_TOP_P)  # 采样生成下一个 token

            out, model_state = model.forward([token], model_state)  # 模型前向传播
            model_tokens += [token] 
            out_tokens += [token]  # 将新生成的 token 添加到输出的 token 列表中

            for xxx in occurrence:
                occurrence[xxx] *= GEN_penalty_decay  # 应用衰减重复惩罚
            occurrence[token] = 1 + (occurrence[token] if token in occurrence else 0)  # 更新 token 的出现次数

            tmp = pipeline.decode(out_tokens[out_last:])  # 将最新生成的 token 解码成文本
            if ("\ufffd" not in tmp) and (not tmp.endswith("\n")):  # 当生成的文本是有效 UTF-8 字符串且不以换行符结尾时
                print(tmp, end="", flush=True) #实时打印解码得到的文本
                out_last = i + 1 #更新输出位置变量 out_last 

            if "\n\n" in tmp:  # 如果生成的文本包含双换行符，表示模型的响应已结束（可以将 \n\n 改成其他停止词）
                print(tmp, end="", flush=True) # 实时打印解码得到的文本
                break #结束本轮推理
    else:
        print("!!! Error: please say something !!!")  # 如果用户没有输入消息，提示“输入错误，说点啥吧！”