一杯茶时间，带你用 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_lightning

一、我们要做怎样的功能

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

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

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

二、代码实现过程

2.1 项目结构

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

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

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

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

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

2.2 激活与补全触发

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

实现 vscode.CompletionItemProvider 的 provideCompletionItems(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 还不够，例如用户输入 a、b、c 时也会触发，容易导致敲一个字母就发一次请求。因此再加一层文档变更的过滤，用 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），用户停一会儿才发请求，减少连打时的无效调用。还可以加两条限制，一是只处理当前活动编辑器的文档，避免在切文件、分屏时误触发，二是与上一次触发至少间隔几百毫秒，进一步避免短时间内重复弹补全。整体触发链路如下：

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 个补全并可取消。五步关系如下：

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

例如下面这样。

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 里 contents 填 Array(numChoices).fill(prefix)，即同一段 prefix 复制多份，利用后端批量接口一次推理出多条不同采样结果；再配上 stream: false、password、max_tokens、temperature、top_p、stop_tokens 等。返回的 data.choices 里，每条取 choice.message?.content || choice.text，trim 掉首尾空白并滤掉空串，得到 string[]。

FIM 补全走 callFIMAPI，请求 POST {baseUrl}/FIM/v1/batch-FIM。prefix、suffix 各为长度为 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: true、retainContextWhenHidden: true 以便跑脚本和在切走时保留状态。panel.webview.html 由 getWebviewContent(completions, languageId) 生成。在打开面板之前，必须把当时的 document 和 position 存到闭包或变量里，因为 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，completions 在 getWebviewContent 里用 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，让光标后的新内容再触发一轮补全，形成补全、选一、再补全的连贯体验。从拿到结果到插入再触发的流程如下：

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

三、如何接入后端

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

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

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:8000 或 http://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	防抖延迟（毫秒）	150–300

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 或弹窗报错，检查 baseUrl、password、端口与防火墙。

---

四、常见问题

为何不能在 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 库在启动前检查并下载模型文件。