Engram：TP 分片与同机 DP 共享的源码研究

Engram 把容量放进巨大的静态查找表，输入 token 的 n-gram 决定要读哪些 rows。部署时的关键问题是：这些 rows 在哪里、谁持有一份、谁发起查询，以及取回的结果怎样进入 hidden state。

本页固定 vLLM 0.31.0 db9527a46873454610df6dbedf79a36d6bf1a7f6 的 NVIDIA 路径。模型机制与全局缓存先读 DeepSeek V4.1 Flash；这里只深读 Engram 的存储、并行与预取。研究结论来自源码，没有独立 GPU 实测。

1. 把四类对象分开

静态表可以让多个 replicas 使用同一份物理权重，查询结果仍依赖每个 replica 的 tokens。它不是把多个请求的动态 KV Cache 混成一份。哈希碰撞由模型训练与 gating 处理，也不代表 cache key 可以忽略模型版本。

2. TP 分片有两个层次

2.1 大表按 hash head 拆分

common/engram.py::ParallelEngramEmbedding 先计算 part_n_hash_cols = ceil(n_hash_cols / num_shards)，再按本 rank 的 head 范围累加 bucket sizes，得到 vocab_start_idx、vocab_end_idx 和 part_num_embeddings。加载器从 checkpoint 切对应 rows，FP8 权重与 ue8m0 scales 使用同样的 row 范围。

这是 head-wise sharding：每个 head 的整段 buckets 属于一个 shard，维度 d 不在这里切开。不同 head 的 prime-sized bucket ranges 略有差别，最后一组还可能 padding，所以 $W/T$ 只是近似；精确值应从每 rank 的 part_num_embeddings 计算。

若 shard 数让某些 rank 没有任何 hash head，构造会拒绝该拓扑。源码 TODO 提及 row-wise sharding，不能把它写成这个版本已实现的能力。

2.2 wkv 按投影输出切分，但有 SP 条件

表查询完成后，拼接得到宽度 n_hash_cols × head_dim 的 memory vector，wkv 输出 (hc_mult + 1) × hidden_size：各 hyper-connection 分支的 key 加一个共享 value。

# 固定源码的关键构造参数；省略其他参数。
self.wkv = ColumnParallelLinear(
    n_hash_cols * layout.head_dim,
    self.dim * (self.hc_mult + 1),
    gather_output=True,
    disable_tp=use_sequence_parallel,
)

Engram.embed() 也有相应分支：普通 TP 沿 head 宽度 gather；SP 路径先 gather rank-major rows，再由 _engram_select_rows 重排，保留本 rank 的 token window。只看到 all_gather(dim=0) 不能把通信对象解释成模型权重。

3. 同机 DP 共享怎样建立

DPSharedEngramStorage 管理一个同节点 Engram DP group 对应的物理分配。相同 TP shard 的 DP peers 执行以下启动流程：

1. group leader 检查空闲空间，在 /dev/shm 创建 backing file，并广播路径。
2. 各进程 mmap(..., MAP_SHARED) 同一文件，从同一 storage 切出 FP8 weights 与 uint8 scales。
3. 每个进程对自己的虚拟地址调用 cudaHostRegister，取得 pinned host memory 的 UVA view。
4. 确认 peers 均完成映射，再释放文件路径；映射继续保持数据存活。
5. leader 写入对应 checkpoint shard，其他 ranks 等待权重加载 barrier。

flowchart TB
    subgraph HOST[同一节点的共享物理主存]
        H0[TP shard 0：rows 与 scales]
        H1[TP shard 1：rows 与 scales]
    end
    H0 --> A0[DP0 / TP0：本请求 rows → GPU staging]
    H1 --> A1[DP0 / TP1：本请求 rows → GPU staging]
    H0 --> B0[DP1 / TP0：另一个请求 rows → GPU staging]
    H1 --> B1[DP1 / TP1：另一个请求 rows → GPU staging]
    A0 <--> A1
    B0 <--> B1

共享发生在同一并行运行组织出的 EDP group 中；同节点、同 IPC namespace 是必要条件，不意味着两个独立启动且未协调的 Engine 自动复用同一 backing file。

4. DP 共享与跨 DP 分片的代价不同

4.1 共享 host 表：各 replica 查自己的 tokens

共享模式将 embedding 的有效 dp_size 设为 1。每个 replica 的 TP ranks 共同覆盖全部 heads，各 TP rank 直接查询自己映射的 host shard；对应的 DP peers 映射同一份物理 shard。gather_engram_hashes(..., dp_shared_memory=True) 返回自己的 hash IDs，不执行 DP hash gather；_ready_rows() 也跳过跨 DP rows exchange。

它消除的是 Engram 查询链的逐步 DP collectives。TP gather、wkv 的 TP 输出 gather、整个模型的 EP/DP 同步以及启动时的 CPU group collectives仍然需要分别计量。

4.2 跨 DP 分片：用通信换更小的每 rank 表

分片模式下，NVIDIA embedding 使用 TP × node-local EDP size 个 shards。一次查询先把同组各 replica 的 hash IDs padding 到公共 token slot，再做 DP all-gather；各 rank 读自己拥有的 heads，随后交换 rows，并选择回本 replica 的 tokens，最后完成 TP 侧的 heads gather。

公共 slot 来自节点内 group 的最大 token count。一个 replica 有长 Prefill、其他 replicas 只有短 Decode 时，padding 与 collective 频率可能成为额外成本。DEAD_ID padding 保障无效位置不被当成合法 row 0。

4.3 配置字段与实际构造路径需要对照

EngramConfig 文档把 embedding_across_dp=False 描述为每 DP replica 独立 TP-sharded 表。但本次 NVIDIA 源码中，ParallelEngramEmbedding.__init__ 直接读取 get_engram_dp_size()，_get_shard_info() 再按有效 dp_size 决定分片；该构造没有直接读取 embedding_across_dp。关闭共享或预检查回退后，不能只凭这个配置字段就认定取得了独立副本基线。

这是一个源码与配置语义需要运行复核的点，不是已经在 GPU 上复现的缺陷。部署记录应增加实际 embed_tokens.dp_size、head/row ranges、EDP group membership 与日志，再用它们标注实验模式。

5. 容量账：省主存与省带宽分别验收

设第 $l$ 个表有 $R_l$ rows，每 row 有 $d$ 个 FP8 字节和 $d/32$ 个 scale 字节，则：

只考虑一节点的完整同机 DP replica group，忽略 head 不均衡：

例如假设 $W=200$ GiB、$T=2$、$D=4$，独立表基线节点约占 800 GiB，共享后约 200 GiB。这个数字只用于展示数量关系，不是 DeepSeek checkpoint 的精确容量或实测结果；多节点仍按各节点的实际 group 数落账。

共享映射可能被各进程 RSS 重复统计，应使用节点物理内存、/dev/shm allocation 和 /proc/<pid>/smaps_rollup 的 PSS/共享页信息交叉检查。只把 RSS 相加会误判共享收益。

若各 replica 都要查询 $N_i$ 个 tokens，每 token 读 $H$ 个 heads，逻辑查询 payload 仍约为：

重复 row、CPU cache、page 粒度与事务放大会改变实际总线流量。共享不会自动把多个 GPU 的访问合并成一次传输；同机 replicas 反而可能争用同一 NUMA memory controller 或 PCIe root。

6. 预取的关键路径与适用边界

model.py 在进入 decoder layers 前生成 hash IDs，并为本地 Engram layers 提交 lookup。各层共享一条 prefetch stream，串行安排 offloaded 查询；消费层只等待自己的 completion event。lookup 的 persistent grid 在 background 模式限制 SM 数，给主计算留出资源。

token IDs → hash → side-stream UVA lookup → staged BF16 rows
                    │ 与前置层计算重叠           │
main-stream 前置层 ───────────────────────→ wait_event
                                             → TP gather → wkv → gate

暴露到关键路径的等待可用 $\max(0,T_{\mathrm{lookup}}-T_{\mathrm{lead}})$ 作为初始假设，但不要把它当作完整延迟公式：查表自身消耗 SM/HBM/host 带宽，会拖慢本来用于掩盖它的前置计算。优化应同时观察 lookup、前置层与端到端 ITL。

配置边界来自源码：共享要求 CPU offload、有效节点内 group、共享 IPC 和足够 /dev/shm；不支持 elastic EP。THP 与共享模式不兼容。DBO/microbatching 会在 V4.1 Engram 配置校验中被拒绝。拓扑/容量预检查可以回退，实际共享分配或注册失败则可能报错，不能统称为静默 fallback。

7. 下一轮研究的最小验证矩阵

先固定 checkpoint、vLLM SHA、GPU/NUMA topology、TP/SP、精度与请求回放；选择能容纳模型的最小同节点配置，再做以下对照。

使用无 profiler 的测量给出性能结果，Trace 负责归因。不要把 event wait 的完整时长都算作可回收收益；先确认同期主 GPU stream 在做什么。

官方 test_engram_dp_shard.py 提供多 rank 表共享、权重加载、graph replay、uneven token padding 与节点边界测试，是未来 GPU 实验的重要入口。本轮阅读了这些测试合同，未执行其 CUDA 多 GPU 用例。

相关页面与源码入口

• Engram NUMA 放置与预取 — 双路 x86 的共享页归属、UVA lookup、预取 FIFO 与受控实验。
• DeepSeek V4.1 Flash — 模型、参考实现与 Runtime 的三个证据范围。
• 分层 KV Cache — 动态状态的加载与恢复；与静态 Engram 表分开落账。
• P/D 分离 — state handoff 与网络关键路径。
• 推理并行 — TP、DP、EP 与拓扑作用域。
• Engram common 源码 — head ranges、wkv 与 SP。
• NVIDIA Engram 源码 — shared storage、UVA、DP exchange 与预取。
• Engram 配置 与 parallel groups — 默认值、校验和节点边界。