阅读视图

发现新文章,点击刷新页面。

一杯茶时间,带你用 RWKV 并发模型做 VS Code 多候选代码补全插件 🤔🤔🤔

在写这份实现教程之前,我已经把该插件的一个版本发布到了 VS Code 扩展市场,在市场中搜索 rwkv 即可找到,你可以先安装试用,再决定是否跟着下文从零实现一版。

本文以这款基于 RWKV 模型的智能代码补全插件为例,讲解从零实现 VS Code 扩展的思路与步骤,并说明如何接入 rwkv_lightning 后端。

该插件通过并发推理一次返回多个不同的补全答案供选择,在侧边栏展示,方便在多种写法之间对比、挑选后再插入,适合写一半、让模型多想几种实现的编码方式;光标后有代码时自动走 FIM(Fill-in-the-Middle)接口做中间填充,否则走普通续写。全文按功能目标、代码实现(项目结构、补全触发、API 调用、Webview 展示与插入)、后端接入组织,后端部分包含硬件要求、模型准备、与 Albatross 的关系、启动服务、模型加载机制、HTTP API、快速测试以及插件配置与验证,文末附常见问题。

下图为在编辑器中触发补全后,并发推理得到的多个不同答案在侧边栏展示、点击即可插入到光标位置的情形。

rwkv-code-completion 效果

前端项目地址:rwkv-code-completion

后端项目地址:rwkv_lightning

一、我们要做怎样的功能

动手写代码之前,首先要考虑我们要实现一个什么样的 VS Code 插件,这决定了后续的架构与实现方式。

在本例中,我们想做一款智能代码补全插件,并事先想清楚四件事。补全结果通过并发推理一次返回多个不同的答案,在侧边栏展示供用户选择,点选后插入。根据光标后是否已有代码,在 FIM(Fill-in-the-Middle)与普通续写之间自动切换接口。在空格、换行、删除等操作时自动触发,并做好防抖与取消,避免频繁请求。服务地址、密码、生成长度、采样参数(temperaturetop_p)、候选数量、防抖延迟等通过 VS Code 设置暴露。把这四件事的对应关系梳理出来,大致如下:

20260121175047

把这些想清楚之后,再按代码实现过程和如何接入后端两部分往下做。

二、代码实现过程

2.1 项目结构

yo code 或手工 scaffold 一个扩展后,核心只需两个源码文件,职责分开,与 VS Code 打交道的放一边,与后端 HTTP 打交道的放另一边,方便维护和单测。

  • src/extension.ts 作为插件入口,在 activate 里实现 CompletionItemProvider、注册补全、用 onDidChangeTextDocument 监听编辑并按条件触发补全;拿到候选列表后,不再往原生 suggest 里塞,而是创建 Webview、渲染多条结果,并处理用户点击插入与插完再补全。
  • src/completionService.ts 负责补全服务,根据有无 suffix 选择调用普通续写接口或 FIM 接口,组装请求体、发 fetch、解析 data.choicesstring[],并透传 AbortSignal 以支持取消。

两者与后端的关系可以概括为:

20260121175344

package.json 里,main 指向打包后的入口(如 ./dist/extension.js),VS Code 按它加载扩展;activationEvents 可设为 onStartupFinished,这样只在 IDE 就绪后才激活,避免启动时卡顿;contributes.configuration 声明 enabledbaseUrlpasswordmaxTokenstemperaturetopPnumChoicesdebounceDelay 等,用户改设置后可通过 vscode.workspace.getConfiguration("rwkv-code-completion") 读到。

构建可用 esbuild 或 tsc,把 extension.ts 等打出到 dist,调试和发布都从 dist 走。

2.2 激活与补全触发

激活时在 activate(context) 里完成两件事,一是向 VS Code 注册谁在什么情况下提供补全,二是监听文档变更,在特定编辑动作后自动调出补全,用户不必每次手动按 Ctrl+Space。

实现 vscode.CompletionItemProviderprovideCompletionItems(document, position, token, context),再用 vscode.languages.registerCompletionItemProvider 挂上去。selector{ pattern: "**" } 表示对所有语言生效;第三参数 triggerChars 是一串字符,当用户输入或删除其中某一个时,VS Code 会来调 provideCompletionItems。这里把空格、换行以及 ASCII 33–126(常见可打印字符)都放进去了,这样在写代码、加空格、换行时都有机会触发,例如:

const selector = { pattern: "**" };
const triggerChars = [
  " ",
  "\n",
  ...Array.from({ length: 94 }, (_, i) => String.fromCharCode(i + 33)),
];
vscode.languages.registerCompletionItemProvider(
  selector,
  provider,
  ...triggerChars,
);

光有 triggerChars 还不够,例如用户输入 abc 时也会触发,容易导致敲一个字母就发一次请求。因此再加一层文档变更的过滤,用 vscode.workspace.onDidChangeTextDocument 监听,只有在本次编辑是删除、换行或输入一个空格时,才在防抖后执行 editor.action.triggerSuggest,从而间接调用 provideCompletionItems。这样可以把触发收敛到更自然的断句、换行场景,例如:

const shouldTrigger = event.contentChanges.some((change) => {
  const isDelete = change.rangeLength > 0 && change.text === "";
  const isNewline = change.text === "\n" || change.text === "\r\n";
  const isSpace = change.text === " ";
  return isDelete || isNewline || isSpace;
});
if (shouldTrigger) {
  debounceTimer = setTimeout(() => {
    vscode.commands.executeCommand("editor.action.triggerSuggest");
  }, config.debounceDelay);
}

防抖时间用 config.debounceDelay(如 150–300ms),用户停一会儿才发请求,减少连打时的无效调用。还可以加两条限制,一是只处理当前活动编辑器的文档,避免在切文件、分屏时误触发,二是与上一次触发至少间隔几百毫秒,进一步避免短时间内重复弹补全。整体触发链路如下:

20260121175403

2.3 补全逻辑与 API 调用

provideCompletionItems 被调用后,先做一轮要不要真的发请求的过滤和节流,再取上下文、调后端、拿 string[]

流程可以拆成五步。一,读配置,若 enabled 为 false 直接 return null。二,防抖,用 setTimeout(..., debounceDelay) 把实际请求放到回调里;若在等待期间又有新的触发,则 clearTimeout 掉上一次,只保留最后一次,这样连续输入时只会发一次请求。三,若此前已有进行中的 fetch,用 AbortController.abort() 取消,再 new AbortController() 给本次请求用。四,取上下文,前缀 prefix 为从文档开头到光标前的文本,document.getText(new vscode.Range(0, 0, position)),过长时截断到约 2000 字符,避免超过后端限制;后缀 suffix 为从光标到往后若干行(如 10 行),主要用来判断光标后是否还有代码,从而决定走 FIM 还是普通续写。五,调用 CompletionService.getCompletion(prefix, suffix, languageId, config, abortController.signal),在 withProgress 里展示正在生成 N 个补全并可取消。五步关系如下:

20260121175421

CompletionService.getCompletion 内部按 suffix 是否非空分支,有后缀则认为用户在中间写代码,走 FIM,否则走普通续写。接口选择如下:

20260121175704

例如下面这样。

async getCompletion(prefix, suffix, languageId, config, signal): Promise<string[]> {
  const hasSuffix = suffix && suffix.trim().length > 0;
  return hasSuffix
    ? this.callFIMAPI(prefix, suffix, config, signal)
    : this.callCompletionAPI(prefix, config, signal);
}

普通补全走 callCompletionAPI,请求 POST {baseUrl}/v2/chat/completions。body 里 contentsArray(numChoices).fill(prefix),即同一段 prefix 复制多份,利用后端批量接口一次推理出多条不同采样结果;再配上 stream: falsepasswordmax_tokenstemperaturetop_pstop_tokens 等。返回的 data.choices 里,每条取 choice.message?.content || choice.text,trim 掉首尾空白并滤掉空串,得到 string[]

FIM 补全走 callFIMAPI,请求 POST {baseUrl}/FIM/v1/batch-FIMprefixsuffix 各为长度为 4 的数组(同一 prefix、同一 suffix 各复制 4 份),对应 4 条并发中间填充;其它参数与普通补全类似,解析方式相同。两处都把 signal 传给 fetch,这样在用户点击取消、或防抖导致下一次触发而 abort() 时,正在进行的请求会被中断,不把过时结果再展示出来。

2.4 Webview 展示与插入

拿到 string[] 之后,不转成 CompletionItem[] 通过 resolve(items) 塞给原生 suggest,因为原生列表单条、偏短,且没法做多列、点击选一等自定义交互。这里改为 resolve(null) 表示不往建议列表里填,同时在 withProgress 里调 showCompletionWebview(document, position, completions, languageId),用 Webview 在侧边栏展示多条候选,支持多选一、点即插、插完再补。

vscode.window.createWebviewPanel 创建 Webview,指定 id、标题、ViewColumn.Two 在侧边打开,以及 enableScripts: trueretainContextWhenHidden: true 以便跑脚本和在切走时保留状态。panel.webview.htmlgetWebviewContent(completions, languageId) 生成。在打开面板之前,必须把当时的 documentposition 存到闭包或变量里,因为 Webview 是异步的,用户可能切文件、移光标,等到点击插入时要以当初触发补全的那次位置为准,否则会插错地方。

const panel = vscode.window.createWebviewPanel(
  "rwkvCompletion",
  "RWKV 代码补全 (N 个选项)",
  vscode.ViewColumn.Two,
  { enableScripts: true, retainContextWhenHidden: true },
);
panel.webview.html = getWebviewContent(completions, languageId);

HTML 里顶部放标题与简短说明,下面一个 div 容器,用 grid-template-columns: 1fr 1fr 做多列布局,每个格子一个 div.code-block,含小标题(序号、字符数、行数)和 <pre><code> 放补全内容。补全文本要先做 HTML 转义再插入,避免 XSS;颜色、背景用 var(--vscode-editor-background) 等,跟主题一致;:hover.selected 给一点高亮,点的时候有反馈。

前端通过 acquireVsCodeApi() 拿到和扩展通信的 API,completionsgetWebviewContent 里用 JSON 注入到页面。每个 code-block 点击时执行 vscode.postMessage({ command: 'insert', code: completions[index] })。扩展侧在 panel.webview.onDidReceiveMessage 里监听,若 message.command === 'insert',先 vscode.window.showTextDocument(targetDocument, ViewColumn.One) 把原文档激活到主编辑区,再用 editor.edit(eb => eb.insert(targetPosition, message.code)) 在事先存好的 targetPosition 插入;插入成功后 panel.dispose() 关掉 Webview,并 setTimeout(..., 300) 后执行 editor.action.triggerSuggest,让光标后的新内容再触发一轮补全,形成补全、选一、再补全的连贯体验。从拿到结果到插入再触发的流程如下:

20260121175751

原生 suggest 只能一条条、样式固定,没法同时展示多条并发结果和自定义交互;用 Webview 可以自己布局、自己处理点击和插入,更适合并发推理、多答案选一的用法。

三、如何接入后端

插件通过 HTTP 调用 rwkv_lightning,需要先部署后端,再在 VS Code 里填好配置。扩展详情页会标注后端部署与配置说明,便于快速上手,下图为扩展市场中的页面示意。

RWKV 代码补全 - 扩展市场页面

接入后端的整体步骤如下。

20260121175818

3.1 硬件要求

重要提示:本后端必须使用 GPU 加速,不支持纯 CPU 运行。

rwkv_lightning 依赖自定义的 CUDA 或 HIP 内核进行高性能推理,因此需要以下硬件之一:

  • NVIDIA GPU:需要支持 CUDA 的 NVIDIA 显卡,并安装 CUDA 工具包
  • AMD GPU:需要支持 ROCm 的 AMD 显卡,并安装 ROCm 运行时

如果您只有 CPU 环境,请使用 llama.cpp 进行 RWKV 模型的 CPU 推理,该项目针对 CPU 进行了专门优化。

3.2 模型文件准备

rwkv_lightning 当前不提供自动下载功能,需要您自行准备模型权重文件。

下载模型权重

RWKV-7 模型的官方权重托管在 Hugging Face 上,推荐从 BlinkDL/rwkv7-g1 仓库下载。模型文件格式为 .pth,例如 rwkv7-g1b-1.5b-20251202-ctx8192.pth

您可以通过以下方式下载:

方式一:使用 huggingface-cli(推荐)

# 首先登录 Hugging Face(如未登录)
huggingface-cli login

# 下载模型文件
huggingface-cli download BlinkDL/rwkv7-g1 \
  rwkv7-g1b-1.5b-20251202-ctx8192.pth \
  --local-dir /path/to/models \
  --local-dir-use-symlinks False

方式二:使用 Python 脚本

from huggingface_hub import hf_hub_download

model_path = hf_hub_download(
    repo_id="BlinkDL/rwkv7-g1",
    filename="rwkv7-g1b-1.5b-20251202-ctx8192.pth",
    local_dir="/path/to/models"
)
print(f"模型已下载到: {model_path}")

路径命名规则

启动服务时,--model-path 支持两种写法。写法一:不带后缀,程序会自动补上 .pth,例如:

--model-path /path/to/rwkv7-g1b-1.5b-20251202-ctx8192
# 实际加载: /path/to/rwkv7-g1b-1.5b-20251202-ctx8192.pth

3.3 与 Albatross 的关系

rwkv_lightning 是基于 Albatross 高效推理引擎开发的 HTTP 服务后端。Albatross 是 BlinkDL 开发的高性能 RWKV 推理引擎,专注于底层计算优化和性能基准测试。

Albatross 项目简介

Albatross 是一个独立的开源项目,GitHub 地址:github.com/BlinkDL/Alb… RWKV-7 模型的高效推理实现,包括:

  • 批量推理支持:支持大规模批量处理,在 RTX 5090 上可实现 7B 模型 fp16 bsz960 超过 10000 token/s 的解码速度
  • 性能优化:集成了 CUDA Graph、稀疏 FFN、自定义 CUDA 内核等优化技术
  • 基准测试工具:提供详细的性能基准测试脚本,用于评估不同配置下的推理性能
  • 参考实现:包含完整的模型实现和工具类,可作为开发参考

性能参考数据

根据 Albatross 官方测试结果(RTX 5090,RWKV-7 7.2B fp16):

  • 单样本解码(bsz=1):145+ token/s,使用 CUDA Graph 优化后可达 123+ token/s
  • 批量解码(bsz=960):10250+ token/s
  • Prefill 阶段(bsz=1):11289 token/s
  • 批量解码(bsz=320):5848 token/s,速度恒定且显存占用稳定(RNN 特性)

rwkv_lightning 的定位

rwkv_lightning 在 Albatross 的基础上,专注于提供生产级的 HTTP 推理服务:

  • HTTP API 接口:提供完整的 RESTful API,支持流式和非流式推理
  • 状态管理:实现三级缓存系统(VRAM、RAM、Disk),支持会话状态持久化
  • 连续批处理:动态管理批次,提高 GPU 利用率
  • 多接口支持:提供聊天、翻译、代码补全等多种应用场景的专用接口

如果您需要深入了解底层实现细节、进行性能调优或对比不同优化方案,建议参考 Albatross 项目的源代码和基准测试脚本。Albatross 提供了更底层的实现细节,而 rwkv_lightning 则专注于提供易用的服务化接口。

3.4 启动推理服务

rwkv_lightning 以 Robyn 版本为主,提供密码认证、多接口、状态管理等特性,适合生产环境使用。Robyn 版本功能更全面,支持密码认证、多接口、状态管理等高级特性,适合生产环境使用。

python main_robyn.py --model-path /path/to/model --port 8000 --password rwkv7_7.2b

如果不需要密码保护,可以省略 --password 参数:

python main_robyn.py --model-path /path/to/model --port 8000

3.5 模型加载机制

了解模型加载机制有助于排查问题和优化性能。

权重加载流程

模型类 RWKV_x070 在初始化时会执行以下步骤:

  1. 读取权重文件:使用 torch.load(args.MODEL_NAME + '.pth', map_location='cpu') 将权重加载到 CPU 内存
  2. 数据类型转换:将权重转换为半精度(dtype=torch.half)以节省显存
  3. 设备迁移:根据硬件平台将权重移动到 GPU
    • NVIDIA GPU:使用 device="cuda"
    • AMD GPU:使用 ROCm 的 HIP 运行时

词表加载

词表文件 rwkv_batch/rwkv_vocab_v20230424.txt 通过 TRIE_TOKENIZER 类自动加载。TRIE 数据结构提供了高效的 token 查找和编码、解码功能。

CUDA、HIP 内核编译

项目包含自定义的 CUDA(NVIDIA)和 HIP(AMD)内核,用于加速 RWKV 的核心计算。这些内核在首次导入 rwkv_batch.rwkv7 模块时通过 torch.utils.cpp_extension.load 自动编译和加载:

  • CUDA 内核:rwkv_batch/cuda/rwkv7_state_fwd_fp16.cu
  • HIP 内核:rwkv_batch/hip/rwkv7_state_fwd_fp16.hip

首次运行时会进行编译,可能需要几分钟时间。编译后的内核会被缓存,后续启动会更快。

3.6 HTTP API 接口

rwkv_lightning 提供了丰富的 HTTP API 接口,支持多种推理场景。

聊天完成接口

  • v1/chat/completions:基础批量同步处理接口,支持流式和非流式输出。
  • v2/chat/completions:连续批处理接口,动态管理批次以提高 GPU 利用率,适合高并发场景。
  • v3/chat/completions:异步批处理接口,使用 CUDA Graph 优化(batch_size=1 时),提供最低延迟。

Fill-in-the-Middle 接口

FIM/v1/batch-FIM:支持代码和文本的中间填充补全,适用于代码补全、文本编辑等场景。

批量翻译接口

translate/v1/batch-translate:批量翻译接口,兼容沉浸式翻译插件的 API 格式,支持多语言互译。

会话状态管理接口

state/chat/completions:支持会话状态缓存的对话接口,实现多轮对话的上下文保持。状态采用三级缓存设计:

  • L1 缓存:VRAM(显存),最快访问
  • L2 缓存:RAM(内存),中等速度
  • L3 缓存:SQLite 数据库(磁盘),持久化存储

流式推理示例

以下示例展示如何使用 v2 接口进行批量流式推理:

curl -N -X POST http://localhost:8000/v2/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "contents": [
      "English: After a blissful two weeks, Jane encounters Rochester in the gardens.\n\nChinese:",
      "English: That night, a bolt of lightning splits the same chestnut tree.\n\nChinese:"
    ],
    "max_tokens": 1024,
    "stop_tokens": [0, 261, 24281],
    "temperature": 0.8,
    "top_k": 50,
    "top_p": 0.6,
    "alpha_presence": 1.0,
    "alpha_frequency": 0.1,
    "alpha_decay": 0.99,
    "stream": true,
    "chunk_size": 128,
    "password": "rwkv7_7.2b"
  }'

3.7 快速测试与性能评估

快速测试

项目提供了测试脚本,可以快速验证服务是否正常运行:

bash ./test_curl.sh

该脚本会发送示例请求到本地服务,检查各个接口的基本功能。

性能基准测试

使用 benchmark.py 脚本可以评估模型的推理性能,包括吞吐量、延迟等指标:

# 需要先修改 benchmark.py 中的模型路径
python benchmark.py

基准测试会输出详细的性能报告,帮助您了解模型在不同配置下的表现。

3.8 插件配置

在 VS Code 中打开设置(可搜索 rwkv-code-completion 或执行命令 RWKV: 打开设置),重点配置:

配置项 说明 示例
rwkv-code-completion.enabled 是否启用补全 true
rwkv-code-completion.baseUrl 后端基础地址,不含路径 http://192.168.0.157:8000http://localhost:8000
rwkv-code-completion.password --password 一致 rwkv7_7.2b
rwkv-code-completion.maxTokens 单次生成最大 token 数 200
rwkv-code-completion.numChoices 普通补全的候选数量(1–50) 24
rwkv-code-completion.debounceDelay 防抖延迟(毫秒) 150300

baseUrl 只需填 http(s)://host:port,插件内部会拼上 /v2/chat/completions/FIM/v1/batch-FIM。若设置界面中仅有 endpoint 等项,可在 settings.json 中手动添加 "rwkv-code-completion.baseUrl": "http://<主机>:<端口>"

3.9 验证接入

可先用 curl -X POST http://<host>:<port>/v2/chat/completions -H "Content-Type: application/json" -d '{"contents":["你好"],"max_tokens":10,"password":"<你的password>"}' 或运行 ./test_curl.sh 确认 v2 与 FIM 接口正常。在任意代码文件中输入、换行或删除,防抖后应出现「🤖 RWKV 正在生成 N 个代码补全...」并弹出侧边栏展示多个候选;若失败,可查看「输出」中该扩展的 channel 或弹窗报错,检查 baseUrlpassword、端口与防火墙。


四、常见问题

为何不能在 CPU 上运行?

rwkv_lightning 的核心计算依赖自定义的 CUDA、HIP 内核,这些内核专门为 GPU 并行计算设计。CPU 无法执行这些内核代码,因此必须使用 GPU。如果您需要在 CPU 上运行 RWKV 模型,请使用 llama.cpp,它提供了针对 CPU 优化的实现。

模型权重应该放在哪里?

模型权重可以放在任何可访问的路径。启动服务时通过 --model-path 参数指定路径即可。路径可以是绝对路径或相对路径,程序会自动处理 .pth 后缀的添加。

首次启动为什么很慢?

首次启动时会编译 CUDA、HIP 内核,这个过程可能需要几分钟。编译后的内核会被缓存,后续启动会快很多。如果希望进一步优化性能,可以考虑使用 torch.compile 模式(详见 README.md 中的 Tips 部分)。

如何选择合适的接口?

  • v1:适合简单的批量推理需求
  • v2:适合高并发场景,需要动态批处理
  • v3:适合单请求低延迟场景(batch_size=1)
  • FIM:适合代码补全和文本编辑
  • state:适合需要保持上下文的对话场景

本插件已按「无 suffix 用 v2、有 suffix 用 FIM」自动选择。

如何实现自动下载模型?

当前版本不提供内置的自动下载功能。您可以在启动脚本中添加下载逻辑,使用 huggingface_hub 库在启动前检查并下载模型文件。

❌