本文基于 CANN ops-nn 仓库中的 Embedding 和 Gather 算子，解析其在 AIGC 文本生成（如 GPT、LLaMA）中的核心作用。
一、文本生成与 Embedding 算子
1.1 AIGC 文本生成的"词典"：Embedding
当 ChatGPT 生成"你好"这两个字时，它并不是直接处理汉字，而是先将文字转换为数字向量。这个转换过程就是 Embedding（词嵌入）。
Embedding 是 AIGC 文本生成的"入口"和"出口"：
- 入口：将用户输入的文字转换为模型能理解的向量
- 出口：将模型输出的向量转换回文字
以 LLaMA-7B 为例：
- 词表大小：32,000 个 Token
- 嵌入维度：4,096
- Embedding 表大小：32,000 × 4,096 = 128M 参数
每次生成一个 Token，都需要从这个巨大的 Embedding 表中查找对应的向量。
Token ID\n你好
向量表示\n[4096]
LM Head\n输出概率
CANN ops-nn 仓库提供了高效的 Embedding 和 Gather 算子，支持大词表场景下的快速查表操作。
1.2 ops-nn 相关算子
| 算子 | 功能 | AIGC 场景 |
|---|---|---|
| Embedding | 词嵌入查表 | Token → 向量 |
| Gather | 通用索引取值 | KV Cache 索引 |
| GatherNd | 多维索引 | Beam Search |
二、ops-nn Embedding 实现
2.1 高效查表机制
渲染错误: Mermaid 渲染失败: Parse error on line 2: ... A[Token IDs
batch, seq\]\] --\> B\[ -----------------------\^ Expecting 'SQE', 'DOUBLECIRCLEEND', 'PE', '-)', 'STADIUMEND', 'SUBROUTINEEND', 'PIPE', 'CYLINDEREND', 'DIAMOND_STOP', 'TAGEND', 'TRAPEND', 'INVTRAPEND', 'UNICODE_TEXT', 'TEXT', 'TAGSTART', got 'SQS'
ops-nn 优化了大词表场景的内存访问模式：
| 词表大小 | 隐藏维度 | 内存占用 | 查表耗时 |
|--------|------|-------|--------|
| 32000 | 4096 | 256MB | 0.05ms |
| 128000 | 4096 | 1GB | 0.12ms |
#### 2.2 位置编码融合
ops-nn 支持 Embedding + 位置编码的融合：
Token Embedding
Add
Position Embedding
输出
融合后减少一次内存读写。
*** ** * ** ***
### 三、Gather 在 AIGC 中的应用
#### 3.1 KV Cache 索引
LLM 推理中，Gather 用于从 KV Cache 中提取历史信息：
输出 KV Cache Token 位置索引 输出 KV Cache Token 位置索引 Gather 操作 提取对应位置的 K/V
#### 3.2 Beam Search 采样
文本生成的 Beam Search 需要 Gather 重排序列：
Beam 候选
5 个序列
计算得分
Top-K 选择
Gather 重排
选择最优序列
继续生成
*** ** * ** ***
### 四、性能优化
#### 4.1 向量化访存
ops-nn 使用向量化指令优化 Gather 的内存访问：
| 优化技术 | 效果 |
|--------|------------|
| 连续访问合并 | 带宽利用率 +50% |
| 预取优化 | 延迟隐藏 |
| 缓存友好布局 | 命中率提升 |
#### 4.2 性能数据
| 操作 | Shape | 耗时 |
|-----------|---------------------------------|--------|
| Embedding | \[1, 2048\] → \[1, 2048, 4096\] | 0.8ms |
| Gather | \[1, 1024, 128\] | 0.15ms |
*** ** * ** ***
### 五、开发者实践
```cpp
// ops-nn Embedding 调用
aclnnEmbedding(workspace, workspaceSize,
weight, indices, output, stream);
// ops-nn Gather 调用
aclnnGather(workspace, workspaceSize,
input, dim, index, output, stream);
```
*** ** * ** ***
### 六、文本生成技术演进
#### 6.1 从 RNN 到 Transformer
文本生成技术经历了重大变革：
| 时代 | 模型 | 特点 | Embedding 使用 |
|------|-------------|-------|--------------|
| 2013 | Word2Vec | 静态词向量 | 预训练 |
| 2017 | Transformer | 注意力机制 | 可学习 |
| 2018 | GPT | 自回归生成 | 大词表 |
| 2020 | GPT-3 | 大规模 | 超大词表 |
| 2023 | LLaMA | 开源 | 32K 词表 |
#### 6.2 Embedding 的重要性
Embedding 作用
离散到连续
语义表示
参数共享
Token ID → 向量
相似词相近
输入输出共享
*** ** * ** ***
### 七、ops-nn Embedding 优化技术
#### 7.1 大词表挑战
| 词表大小 | 嵌入维度 | 参数量 | 内存占用 |
|------|------|------|-------|
| 32K | 4096 | 128M | 256MB |
| 128K | 4096 | 512M | 1GB |
| 256K | 4096 | 1B | 2GB |
#### 7.2 访存优化
是
否
Token IDs
连续访问?
合并访存
随机访存
高效
优化: 预取 + 缓存
*** ** * ** ***
### 八、Gather 在 LLM 中的应用
#### 8.1 KV Cache 索引
位置索引
Gather
KV Cache
选中的 K/V
#### 8.2 Beam Search 重排
5 个候选序列
计算得分
选择 Top-5
Gather 重排
新的 5 个序列
*** ** * ** ***
### 九、AIGC 文本生成应用
#### 9.1 LLM 推理流程
输入 Token IDs
Embedding 查表
Transformer 层 ×N
LM Head
logits
采样
输出 Token
#### 9.2 输入输出 Embedding 共享
| 模型 | 共享方式 | 参数节省 |
|-------|------|------|
| GPT-2 | 共享 | 50% |
| LLaMA | 不共享 | 0% |
| Qwen | 共享 | 50% |
*** ** * ** ***
### 十、性能优化策略
#### 10.1 Embedding 优化
| 优化技术 | 方法 | 收益 |
|------|----------------|--------|
| 量化 | INT8 Embedding | ��存减半 |
| 分片 | 多卡分布 | 支持更大词表 |
| 缓存 | 热门 Token 缓存 | 减少访存 |
#### 10.2 Gather 优化
| 优化技术 | 方法 | 收益 |
|------|---------|------|
| 向量化 | SIMD 并行 | 吞吐提升 |
| 预取 | 提前加载 | 隐藏延迟 |
| 合并 | 连续索引合并 | 减少访存 |
*** ** * ** ***
### 十一、开发者实践指南
#### 11.1 完整调用示例
```cpp
#include "aclnn/acl_nn.h"
// Embedding 查表
aclnnStatus embeddingStatus = aclnnEmbedding(
workspace, workspaceSize,
weight, // [vocab_size, hidden_dim]
indices, // [batch, seq_len]
output, // [batch, seq_len, hidden_dim]
stream
);
// Gather 索引
aclnnStatus gatherStatus = aclnnGather(
workspace, workspaceSize,
input, // [batch, seq_len, dim]
1, // dim
index, // [batch, num_indices]
output, // [batch, num_indices, dim]
stream
);
// GatherNd 多维索引
aclnnStatus gatherNdStatus = aclnnGatherNd(
workspace, workspaceSize,
input, // [batch, seq_len, dim]
indices, // [num_indices, 2] (batch_idx, seq_idx)
output, // [num_indices, dim]
stream
);
// LLM 输入处理
void llmInputProcess(
int* tokenIds, // [batch, seq_len]
aclTensor* output // [batch, seq_len, hidden]
) {
// 1. Token Embedding
aclnnEmbedding(workspace, workspaceSize,
tokenEmbedding, tokenIds,
tokenVectors, stream);
// 2. 位置编码 (如果使用绝对位置)
aclnnEmbedding(workspace, workspaceSize,
positionEmbedding, positionIds,
positionVectors, stream);
// 3. 相加
aclnnAdd(workspace, workspaceSize,
tokenVectors, positionVectors, 1.0,
output, stream);
}
// Beam Search 重排
void beamSearchReorder(
aclTensor* sequences, // [batch, beam, seq_len]
aclTensor* beamIndices, // [batch, beam] 选中的 beam 索引
aclTensor* output
) {
// 使用 Gather 重排序列
aclnnGather(workspace, workspaceSize,
sequences, 1, beamIndices,
output, stream);
}
```
#### 11.2 常见问题与解决方案
| 问题 | 原因 | 解决方案 |
|------|--------------|---------|
| 内存不足 | 词表过大 | 使用量化或分片 |
| 查表慢 | 随机访存 | 优化访存模式 |
| 索引越界 | Token ID 超范围 | 添加边界检查 |
*** ** * ** ***
### 十二、总结与展望
#### 12.1 核心要点
CANN ops-nn 仓库中的 Embedding 和 Gather 算子具有以下特点：
* **大词表支持**：优化的访存模式
* **高效查表**：向量化实现
* **灵活索引**：支持多种 Gather 变体
* **AIGC 适配**：针对 LLM 推理优化
#### 12.2 LLM 部署建议
| 场景 | 推荐配置 | 理由 |
|-------------|--------------|--------|
| 大词表 | 量化 Embedding | 节省内存 |
| Beam Search | 优化 Gather | 提升效率 |
| 长序列 | KV Cache 索引 | 减少重复计算 |
*** ** * ** ***
**相关链接：**
* 🏠 CANN 组织主页：