作者：昇腾实战派 & 2012泊松实验室
知识地图：https://blog.csdn.net/Lumos_Lovegood/article/details/161455142

背景概述

2026年1月12日，DeepSeek发布论文《Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models》，提出Engram条件记忆模块，其作为MoE之外一种高效的稀疏扩展方式，可以复用于其他模型，提升推理效率。

昇腾RecSDK，是专注使能互联网市场搜索推荐广告场景的开发套件，快速完成了在昇腾NPU上对Engram的支持，这篇文章将带您基于RecSDK和昇腾算力体验Engram的推理加速。

RecSDK为Engram训推加速

Engram的设计核心是：针对搜索类任务，用查询代替计算，通过“将海量静态知识转化为可高效检索的Embedding”用于查询，提升了模型性能。昇腾RecSDK中的DynamicEmb，恰好提供了Embedding的多级缓存和预取技术，结合Engram模块，可进一步提升训推性能。RecSDK实现了以下核心优化：

1）多级缓存，提升热点Embedding性能

DRAM、高速片上内存，构建多级缓存：

N-gram符合Zipf分布，我们将高频词放到速度更快的存储介质中，低频词则反之，可提升模型性能。具体来说，我们在NPU高速片上内存中构建了一个热点缓存，而将全量Embedding表存放在Host DRAM中。

热点Embedding实时更新：

为了更准确保持高频词常驻缓存，我们对Embedding做访频统计，并在此基础上实现LFU缓存策略，控制热点数据准入与淘汰均，保持热点缓存实时更新。

2）缓存预取技术，降低时延

Engram可划分成两个阶段：Sparse Retrieval（即Embedding查表）和Context-aware Gating，如下图红框所示：

我们发现，第一阶段Sparse Retrieval只依赖Input Tokens，因此可让这一阶段异步执行，实现Embedding的预取，从而掩盖Embedding查表的时延。

说明：热点缓存和Embedding预取的代码后续会在RecSDK开源仓陆续开源。

快速上手Engram推理

Engram Demo昇腾适配

1）环境准备

pip install torch torch_npu numpy transformers sympy

2）代码下载

git clone https://github.com/deepseek-ai/Engram.git

3）代码适配

在昇腾NPU上运行Engram Demo，需要将模型和数据放置于NPU Device。

device = torch.device('npu')
TransformerBlock(layer_id=layer_id).to(device)

3）代码执行

python engram_demo_v1.py

使用CCI3数据集进行推理

为了测试Engram在开源数据集上的推理表现，我们使用CCI3-HQ数据集。

1. 获取数据集

https://huggingface.co/datasets/BAAI/CCI3-HQ

需要从huggingface获取该数据集的授权

2. 安装依赖包

pip install datasets huggingface_hub

3. 登录hugging_face

hf auth login

4. 使用数据集

我们将demo中的代码做简单修改，使用CCI3-HQ数据集进行推理

class LLMModel(nn.Module):
    def __init__(self, backbone_config, engram_cfg):
        super().__init__()
        self.LLM = [
            nn.Embedding(backbone_config.vocab_size, backbone_config.hidden_size),
            *[TransformerBlock(layer_id=layer_id) for layer_id in range(backbone_config.num_layers)],
            nn.Linear(backbone_config.hidden_size, backbone_config.vocab_size)
        ]
        self.tokenizer = AutoTokenizer.from_pretrained(engram_cfg.tokenizer_name_or_path, trust_remote_code=True)

    def forward(self, text):
        input_ids = self.tokenizer(text, return_tensors='pt').input_ids
        for idx, layer in enumerate(self.LLM):
            if idx == 0:
                hidden_states = self.LLM[0](input_ids)
                ## mock hyper-connection
                hidden_states = hidden_states.unsqueeze(2).expand(-1, -1, backbone_config.hc_mult, -1)
            elif idx == len(self.LLM) - 1:
                ## mock hyper-connection
                hidden_states = hidden_states[:, :, 0, :]
                output = layer(hidden_states)
            else:
                hidden_states = layer(input_ids=input_ids, hidden_states=hidden_states)

        print("✅ Forward Complete!")
        print(f"{input_ids.shape=}\n{output.shape=}")

if __name__ == '__main__':
    llm = LLMModel(backbone_config, engram_cfg)
    from datasets import load_dataset
    dataset = load_dataset("BAAI/CCI3-HQ", streaming=True, split="train")
    dataset = iter(dataset)
    for i in range(100):
        data = next(dataset)
        llm(data['text'])

测试结果

实验共跑了3组性能数据，包括1个对照组，2个实验组。

对照组：使用官方开源代码中的推理模式，引入CCI3-HQ数据集，使用单卡推理，超参使用默认配置（如下）

实验组1：对照组+热点Embedding缓存（缓存比例10%）

实验组2：对照组+热点Embedding缓存+Embedding预取

数据集：CCI3-HQ
序列长度：4096
engram_vocab_size：10342400
n_embed_per_ngram：512
n_head_per_ngram：8
engram_layers：1，15
max_ngram_size：3

1）对照组

2）实验组1

3）实验组2

对比3组数据可见：

1）对照组：推理端到端时延26.72ms，其中Emb Lookup 11.63ms，约占45%。

2）实验组1：推理端到端时延20.94ms，其中Emb Lookup 4.89ms，约占23%。Emb Lookup相比对照组提升2.38x，推理端到端相比对照组提升1.28x。

3）实验组2：推理端到端时延19.22ms，推理端到端相比对照组提升1.39x，其相比实验组1的额外收益来自于计算对Embedding部分预取时间的掩盖。

结语

Engram开源代码链接：https://github.com/deepseek-ai/Engram

RecSDK开源代码链接：https://gitee.com/ascend/RecSDK

CCI3-HQ开源数据集链接：https://huggingface.co/datasets/BAAI/CCI3-HQ