用 burn 算相关文章

主线位置：整段计算是一串算子主线：逐元素运算、按维求和、一次 matmul、一次 topk_with_indices。在 Flex 上每个 op 当场算完，matmul 落到 gemm；在 Metal 上 op 先登记进融合流，读回结果时才编译、调优、提交内核。没有求导主线和训练主线。

flowchart TB
  files["读 docs 下 index.md"] --> tok["分词 计数"]
  tok --> vocab["定词表"]
  vocab --> sparse["稀疏行"]
  sparse --> plainrs["纯 Rust 倒排表"]
  sparse --> dense["稠密 N x V 矩阵"]
  dense --> tfidf["TF-IDF 归一化"]
  tfidf --> mm["matmul 得 N x N"]
  mm --> topk["topk_with_indices"]
  plainrs --> out["每页前 5"]
  topk --> out

本站（开源，仓库 yuxuetr/rustineverything）是 dioxus 全栈应用，不用 burn。这一页回答一个具体问题：给文档页加一个「相关文章」列表，按正文相似度挑出最像的 5 篇，用 burn 算要付出什么，和不用 burn 比又如何。站点源码链接指向提交 8729cfb，只读，没有改动、没有构建站点。

数据是私有内容仓库 docs/ 在提交 dd09e82 时的快照：13 卷、705 个 index.md（本卷当时写到 07）。页面上只给标题和统计。实验 crate 是 exp-burn/site-related：分词和计数是共用的纯 Rust 代码，相似度部分写了三遍，由 feature 选择：不开 feature 时完全不依赖 burn；flex 依赖 burn = { version = "=0.22.1", features = ["flex"] }；metal 在此之上加 burn/metal。全部是 release 构建，平台 M2 Max、macOS 27、rustc 1.99.0。本站的服务器是 Linux、没有 GPU，下文所有数字都只在 macOS 上测过；Linux 上的情况只读了源码和部署文件。测量时机器上同时跑着别的编译任务（构建时 load average 在 100 上下，运行时在 50 上下），绝对时间偏大、波动大，每处都给了当时的负载，比较时看同一时段的相对关系。

本站现在怎么读文档

文档的目录由 content_dir 决定：设了 CONTENT_DIR 就读 $CONTENT_DIR/docs。页面正文由 server function get_doc_content 每次请求从磁盘读一次 index.md，解析 frontmatter；上一篇、下一篇由 doc_neighbors 当场扫这一卷的目录得出。没有任何预先算好的东西。

站点里已经有一份全部正文：搜索模块用 tantivy 建了索引，分词器是 jieba 加小写化（SearchEngine::open_or_create 里注册），标题和正文都存在索引里。文档由 walk_docs 收集：每个目录的 index.md，跳过 _ 和 . 开头的目录。索引持久化在卷里，进程启动后由 main.rs 里的定时任务立即做一次按 mtime 的增量刷新，之后每 10 分钟一次。下面的实验按 walk_docs 的规则取页面。

部署决定了「什么时候算」有哪些选项：

• 镜像在 CI 里构建，服务器上不编译。 release.yml 只出 linux/amd64 镜像，运行阶段是 debian:trixie-slim，只装了 ca-certificates 和 tini：没有 GPU 驱动，也没有 Vulkan 加载器。
• 内容不在镜像里。 生产覆盖文件把宿主上的内容仓库只读挂进容器，CONTENT_DIR=/app/content；更新文章只在宿主上 git pull，不重建镜像。所以「在构建镜像时算」做不到：构建镜像时看不到内容。能选的是在内容仓库一侧离线算好、在进程启动时算、或者每次请求算。
• 内存有上限。 app 容器 mem_limit: 768m，同一台机器上还跑着别的服务。

做法：分词和 TF-IDF

相关度用最朴素的 TF-IDF（词频–逆文档频率，term frequency–inverse document frequency）向量加余弦相似度：两页共有的词越多、这些词在全站越少见，两页越像。

分词。 文档是中文夹着大量代码标识符。本站搜索用的是带词典的 jieba；这里为了让三个版本只在相似度计算上有差别，用一个不需要词典、二十几行就能写完的办法：连续的汉字取相邻两字（「自动」「动微」「微分」），只有一个字时保留单字；ASCII 的字母、数字、下划线连成一个标识符，转小写，至少 2 个字符、不能全是数字；其余字符都是分隔。二元组会产生「动微」这样跨词的片段，但它总是和「自动」「微分」一起出现，只是把同一个信号多数一次，不会让不相关的两页变得相似；真正区分页面的是「融合」「连接池」这类词和 float_matmul、pooloptions（转小写后的 PoolOptions）这类标识符，两者都被完整保留。分词的核心：

  for c in text.chars() {
    if is_han(c) {
      flush_ident(&mut ident, &mut emit);
      han.push(c);
    } else if c.is_ascii_alphanumeric() || c == '_' {
      flush_han(&mut han, &mut emit); // 连续汉字 -> 相邻二元组
      ident.push(c.to_ascii_lowercase());
    } else {
      flush_han(&mut han, &mut emit);
      flush_ident(&mut ident, &mut emit); // >= 2 个字符、不全是数字才算
    }
  }

页面的文本是 frontmatter 里的 title 加正文，正文里 [文字](链接) 只留文字，否则每页都会多出 https、github、blob。代码块原样保留。

加权。 词频取 1 + ln(c)，逆文档频率取 ln((N+1)/(df+1)) + 1，两者相乘后按行做 L2 归一化，归一化后两行的点积就是余弦相似度。

词表。 只出现在一页里的词（df = 1）不可能让两页相似，出现在一半以上页面里的词（df > 352，例如「一个」「主线」「这一」和每页都有的 mermaid、flowchart）几乎不区分页面，两者都去掉。三个版本用同一个词表：

pages=705 chars=8956836 tokens=1773655 distinct_terms=98834
vocab=52120 (dropped df=1: 46550, df>352: 164) nnz=656206 density=1.79% dense_f32=140.2 MiB

705 页、约 896 万字符，切出 177 万个词元、9.9 万个不同的词；词表留下 52 120 个。每页平均只用到其中约 930 个，矩阵的非零元占 1.79%。同样的数据摆成 705 × 52 120 的 f32 稠密矩阵，是 140.2 MiB。

burn 版：一张稠密矩阵

burn 0.22.1 没有稀疏张量（burn-tensor、burn-backend 里没有稀疏存储，仓库历史里也没有过），所以要先把计数摊成稠密矩阵，再交给张量运算：

pub fn related(counts: Tensor<2>, k: usize) -> (Tensor<2>, Tensor<2, Int>) {
  let device = counts.device();
  let [n, _] = counts.dims();
  let present = counts.clone().greater_elem(0.0).float(); // N x V, 0/1
  let df = present.clone().sum_dim(0); // 1 x V
  let idf = ((df + 1.0).recip() * (n as f32 + 1.0)).log() + 1.0;
  let tf = (counts.clamp_min(1.0).log() + 1.0) * present; // 1 + ln c, 0 where c = 0
  let w = tf * idf; // broadcast over rows
  let norm = w.clone().square().sum_dim(1).sqrt().clamp_min(1e-12); // N x 1
  let unit = w / norm;
  let sim = unit.clone().matmul(unit.transpose()); // N x N cosine
  let sim = sim - Tensor::<2>::eye(n, &device) * 2.0; // never pick yourself
  sim.topk_with_indices(k, 1)
}

每一行都是 Tensor 与 Kind 里讲的泛型方法，(1, V) 和 (N, V) 相乘按行广播。设备由调用方传进来的张量决定：Device::flex() 或 Device::metal(DeviceKind::DefaultDevice)，同一段代码不改一个字。分词、计数、摊成稠密矩阵都在 burn 之外，用的是普通的 Vec，最后用 TensorData::new(dense, [n, v]) 装进张量；读回用 into_data()。

同一串 op 在两个后端上走的路不同：

• Flex：每个 op 当场算完、分配新的缓冲（能原地写的逐元素运算会原地写，见 Flex）。matmul 的规模是 705 × 52 120 × 705，远超 192³，交给 gemm 用全部线程。topk_with_indices 没有专门实现，走 float_topk_with_indices 的默认做法：整行排序一次，再各取前 k 个。
• Metal：op 先被 Fusion 记成 IR，逐元素的几步可以融合成一个内核，matmul 交给 cubek，topk_with_indices 由 CubeBackend 的 float_topk_with_indices 用一次归约内核同时求出值和下标；这些都在 into_data 读回时才发生，第一次遇到的形状要编译内核、跑 autotune。

纯 Rust 版：稀疏行加倒排表

不用 burn 时，自然的写法是保留稀疏：每页一个 (词 id, 权重) 的有序列表，再按词建倒排表，对每一页只沿着它自己用到的词累加分数：

  let mut scores = vec![0f32; n];
  weighted
    .iter()
    .enumerate()
    .map(|(d, row)| {
      scores.iter_mut().for_each(|s| *s = 0.0);
      for &(t, w) in row {
        for &(other, w2) in &postings[t as usize] {
          scores[other as usize] += w * w2;
        }
      }
      scores[d] = f32::NEG_INFINITY;
      let mut idx: Vec<usize> = (0..n).collect();
      idx.select_nth_unstable_by(k, |&a, &b| scores[b].total_cmp(&scores[a]));
      // 前 k 个按分数排序后返回

加权和归一化与 burn 版逐项相同，只是不碰零。稠密的 matmul 要做 705 × 705 × 52 120 ≈ 2.6 × 10¹⁰ 次乘加；倒排表只在两页共有某个词时才乘，次数是词表里每个词的 df² 之和：

multiply-adds: dense matmul n*n*V = 2.590e10, inverted index sum(df^2) = 5.288e7

相差约 490 倍。稠密矩阵里 98.2% 是零，gemm 不知道这一点，照样逐个乘。

结果对不对、合不合理

burn 版和纯 Rust 版在同一进程里各算一遍，逐页比较前 5 的下标集合和分数：

agree with plain: same top-5 set on 705/705 pages, max |score diff| = 1.28e-6

Flex 和 Metal 都是 705 页全部一致，分数最大相差 1.28 × 10⁻⁶，是 f32 累加顺序不同带来的。下面是四个页面的前 5（只列标题和分数）：

[burn/03-backend/flex] Flex：CPU 后端
  0.239  Tensor 与 Kind  (burn/02-tensor/tensor-and-kind)
  0.237  burn-cubecl 后端  (burn/06-cubecl/cubecl-backend)
  0.204  后端与分派  (burn/03-backend)
  0.197  Backend trait  (burn/03-backend/backend-trait)
  0.193  Dispatch  (burn/03-backend/dispatch)
[axum/02-handler/state] 状态：State 与 FromRef
  0.360  expected `Router`, found `Router<AppState>`  (axum/10-pitfalls/router-state-mismatch)
  0.351  一个请求穿过 axum  (axum/01-overview/request-through-axum)
  0.268  Router 的内部结构  (axum/04-routing/router-internals)
  0.261  layer 与 route_layer  (axum/05-middleware/layer-and-route-layer)
  0.260  取舍汇总  (axum/09-design/tradeoffs)
[sqlx/04-pool/acquire] 获取连接
  0.390  归还、Drop 与维护任务  (sqlx/04-pool/release-and-maintenance)
  0.353  PoolOptions 与默认值  (sqlx/04-pool/pool-options)
  0.264  本站怎么用 sqlx  (sqlx/09-in-this-site/sqlx-in-site)
  0.251  连接池  (sqlx/04-pool)
  0.249  pool timed out while waiting for an open connection  (sqlx/08-pitfalls/pool-timed-out)
[serde/07-pitfalls/missing-field] missing field `x`
  0.305  生成的 Deserialize  (serde/03-derive/generated-deserialize)
  0.184  value 模块与 IntoDeserializer  (serde/02-core/value-module)
  0.184  错误  (serde/02-core/errors)
  0.181  'invalid type: string "1", expected u32'  (serde/07-pitfalls/invalid-type)
  0.170  Content：缓冲的代价  (serde/04-enums/content)

按人的判断，这些都说得过去：State 那页找到了「Router 状态类型不匹配」这个坑，连接池那页找到了池超时的坑和本站用 sqlx 的页面，missing field 那页第一名是讲生成的 Deserialize 的页面，而缺字段的报错正是在那里产生的。上一篇 / 下一篇只连侧栏里相邻的页，这些跨部分的推荐（坑页面、「在本站中」）原本要读者自己在侧栏里找。

但也能看出它的局限：推荐几乎都在同一卷里。对全部 705 页统计：

locality: same volume 3115/3525 (88.4%), same part 1150/3525 (32.6%), pages with no cross-volume hit 530/705

3525 条推荐里 88.4% 在同一卷，705 页里有 530 页的前 5 一条跨卷的都没有。原因很直接：每卷有自己的一套类型名和函数名（Router、PoolOptions、Dispatch），它们正是 IDF 最看重的词。想要「axum 的 State 页推荐 tower 的 Service 页」这类跨卷关系，TF-IDF 给不了，要靠人工标注或语义向量。

代价

编译

每种配置用一个全新的 target 目录 cargo build --offline --release，依赖源码已在本地缓存。real 是墙钟时间，受当时负载影响很大；user + sys 是 CPU 时间，更能比较：

== flex: cargo build --offline --release --features flex (load: 100.88 107.87 76.57)
    Finished `release` profile [optimized] target(s) in 2m 50s
      170.35 real       405.34 user        17.32 sys
== metal: cargo build --offline --release --features metal (load: 112.62 111.54 88.90)
    Finished `release` profile [optimized] target(s) in 9m 34s
      574.41 real      1814.25 user        58.43 sys

（「编译的包」数的是 cargo 输出里的 Compiling 行，包括本 crate。五次构建时 load average 在 98 到 113 之间。）

• 只开 flex 就要编 160 个包、约 7 分钟 CPU 时间，二进制从 0.53 MiB 变成 8.8 MiB。
• 关掉 burn 的默认 feature 只少编两个包：burn-autodiff 和 burn-optim；burn-nn 是 burn 的非可选依赖，照样在。二进制小 3.5 MiB，编译时间看不出差别。注意要手动把 rayon、simd 加回来：它们是 burn-flex 的默认 feature，经 burn 的默认 feature 才传下去，关掉后 gemm 就只用一个线程。
• 加上 GPU 运行时，包数翻到 416、CPU 时间翻 4 倍多，二进制 100 MiB 以上。metal 和 wgpu 差不多，因为 metal 就是 wgpu 加 MSL 编译器（见 wgpu 与 Metal）。Linux 上的 wgpu 会编进 Vulkan 后端，体积和时间没有测。

运行

site-related <docs 目录> <plain|flex|metal> 5 <示例页面...> 在一个进程里读文件、分词、定词表，然后把相似度计算连做 5 次；burn 版每次都从 TensorData 重新建张量，计时包括上传、计算和 into_data 读回，同一进程里再跑 5 次纯 Rust 版作对照。整个进程套在 /usr/bin/time -l 下，取「maximum resident set size」。五次运行接连进行，load average 在 52 上下。纯 Rust 和 Metal 冷缓存两次的原始输出：

== plain2: .../target-site-clean-plain/release/site-related plain 5 (load: 52.03 74.94 86.02)
load=58.1ms tokenize+count=174.3ms vocab=3.5ms
plain: first=63.6ms steady(mean of 4)=106.4ms min=84.0ms
        0.73 real         0.55 user         0.04 sys
            64602112  maximum resident set size
...
== metal2-cold: .../target-site-clean-metal/release/site-related metal 5 (load: 53.91 74.60 85.77)
load=65.0ms tokenize+count=174.0ms vocab=3.4ms
metal: dense counts built in 8.1ms
metal: first=9556.9ms steady(mean of 4)=44.7ms min=40.0ms
plain (same process): first=57.9ms steady(mean of 4)=97.4ms min=91.4ms
agree with plain: same top-5 set on 705/705 pages, max |score diff| = 1.28e-6
       10.46 real         1.85 user         0.96 sys
           728252416  maximum resident set size
...
          2854274080  peak memory footprint

（Flex 和 Metal 两行的「整个进程」也包括同一进程里 5 次纯 Rust 对照。）

• 时间。 分词计数在三个版本里是同一段代码，约 0.17 s。相似度这一步纯 Rust 约 0.1 s；Flex 慢 5 到 10 倍，原因就是上面那 490 倍的乘加次数，gemm 再快也补不回来。Metal 稳定后约 40 ms，是三者里最快的，但它的「第一次」在冷缓存时是 9.6 s（编内核加 autotune，机制见 autotune），热缓存时仍有 0.5 s（进程内的内核编译不落盘，见 wgpu 与 Metal）。对一个启动时只算一次的任务，Metal 的第一次就是它的全部成本。这次的调优结果写进了 crate 目录下的 target/environment/default.db。
• 内存。 纯 Rust 进程峰值 61.6 MiB；Flex 版 546 MiB，是它的 8.9 倍，约等于 3.9 个 140 MiB 的稠密矩阵：主机上留着一份计数矩阵，每次 from_data 用的是它的拷贝，present、tf、w、unit 也都是 N × V 的张量。Metal 版 RSS 695 MiB，冷缓存那次的「peak memory footprint」到了 2.66 GiB，热缓存时是 1.05 GiB，多出来的部分只出现在跑 autotune 的那次。对照部署：app 容器的上限是 768 MiB，Flex 版这一步就要用掉七成，还没算站点自己。

结论

不为这个功能引入 burn。 在本站的条件下，burn 的三个长处都用不上：没有要训练的参数，用不上自动微分；服务器没有 GPU，用不上 cubecl；计算本身是稀疏的，稠密张量反而是负担。代价却都在：多编 159 个包，二进制大 16 倍，内存大 9 倍，速度还慢。

如果要加「相关文章」：

• 怎么算： 纯 Rust 的稀疏版，或者直接用站点已有 tantivy 索引的 MoreLikeThisQuery（上面的没有蠢问题）。后者不用新依赖、分词和搜索一致，代价是每次查询要在索引上跑一个由最多 25 个词（默认的 max_query_terms）组成的查询。
• 什么时候算： 构建镜像时做不到，内容不在镜像里。每次请求算也不合适：IDF 要用全站的文档频率，单页也得先有全量统计。剩下两种都可行，而且都便宜：一是进程启动时算一次，和搜索索引的增量刷新放在一起，按 mtime 判断内容变了才重算，全站 705 页约 0.3 s、几十 MiB，结果只需要保留 705 × 5 个下标；二是在内容仓库一侧离线算，把结果写成一个 JSON 跟着内容一起 git pull，服务器上零计算，代价是内容仓库多一个生成步骤、忘了跑就过期。前者更符合站点现在「内容只靠 git pull 更新」的做法。
• 如果以后真要用 burn（例如改用一个小模型算语义向量，才能跨卷推荐）：只开 flex，并在代码里写 Device::flex()。不要开 wgpu 或 metal：除了编译和体积，开了之后 Device::default() 会先选 wgpu，在没有 GPU、镜像里也没有 Vulkan 加载器的容器里取不到适配器。这一条只读了源码和部署文件，没有在 Linux 上跑过。default-features = false 只省两个包，可以不折腾。

最后再说一次范围：以上所有时间、内存、二进制大小都是在 macOS（M2 Max）上测的。服务器是 x86_64 的 Linux：gemm 走的是 x86 的代码路径（burn-flex 的 x86-v4 即 AVX-512 默认不开），线程数和内存带宽也不同，绝对数字会变；但稠密和稀疏差 490 倍的乘加次数、140 MiB 的稠密矩阵和随之而来的内存倍数，取决于数据的形状，不取决于平台。

小结

• 本站不用 burn。文档每次请求从磁盘读；搜索用 tantivy 加 jieba，启动时和每 10 分钟增量刷新索引；镜像在 CI 里出 linux/amd64，内容只读挂载、git pull 更新，容器内存上限 768 MiB，没有 GPU。
• 705 页、177 万个词元，字的二元组加标识符分词；去掉 df = 1 和 df > N/2 的词后词表 52 120，非零元 1.79%。
• burn 0.22.1 没有稀疏张量，同一计算要摊成 140 MiB 的稠密矩阵；matmul 要做 2.6 × 10¹⁰ 次乘加，倒排表只要 5.3 × 10⁷ 次。
• 三个版本的前 5 在 705 页上完全一致；推荐大多合理，但 88.4% 在同一卷，530 页没有一条跨卷推荐。
• 代价（macOS 实测）：只开 flex 多编 159 个包、约 7 分钟 CPU 时间，二进制 0.53 → 8.8 MiB，峰值内存 62 → 546 MiB，计算还慢 5 到 10 倍；Metal 稳定后 40 ms，但冷启动 9.6 s、二进制 106 MiB，服务器也没有 GPU。
• 结论：用纯 Rust（或 tantivy 的 MoreLikeThisQuery）在进程启动时、内容变化后算一次；真要用 burn 只开 flex 并显式用 Device::flex()。