只要是自己写过一次推理引擎你才会意识到词嵌入这层有多容易被低估这是《自己动手写LLM推理引擎》的第四章讲词嵌入Token Embedding。前面我们已经把权重格式、张量加载、模型结构骨架捋清楚了这一章处理的是真正进入Transformer之前的第一个算子把token ID变成向量。说实话单看概念词嵌入可能是整个LLM推理引擎里“最没什么好讲”的一层——它就是一个查表操作输入是整数ID输出是固定维度的浮点向量。很多教程一句话带过。但当我真正着手实现的时候发现这一层埋的细节远比想象中多权重在磁盘上的布局、内存里的对齐方式、量化格式的兼容、batch内逐token的取行效率、以及和后续LayerNorm、旋转位置编码之间的衔接任何一个环节没处理好模型跑出来的效果都是灾难性的。这一章适合正在手写推理引擎的人也适合那些已经会用现成框架、但想知道“输入层到底发生了什么”的读者。我会尽量把踩过的坑、对比过的方案、以及最终采用的实现都交代清楚。1. 词嵌入的真实角色不是“查表”这么简单1.1 输入输出与位置推理流水线的第一棒先明确推理引擎的数据流。假设我们已经完成了tokenizer编码拿到一段prompt的token ID序列比如[1, 306, 1128, 338, 263, 8876, 322]其中1通常是BOS句首标记。接下来引擎要做的是把每个ID映射为一个稠密向量向量维度就是模型配置里的hidden_size常见的是4096。这一步就是嵌入层的工作。在PyTorch里它是nn.Embedding(vocab_size, hidden_size)在前向计算时输入input_ids输出[seq_len, hidden_size]的张量。但在推理引擎的实现里我不会真的去做一次矩阵乘法而是直接按ID去权重表里取对应行。原因很简单nn.Embedding的前向在推理时本来就是查表没必要绕道走通用矩阵乘。一个容易忽略的点是嵌入层是唯一一个输入不是浮点张量的算子。前面的tokenizer输出是整数序列嵌入层之后才是浮点向量流。这意味着类型转换、越界检查、显存带宽的利用方式都和其他层不太一样。推理引擎的算子设计里它属于“密集但极简”的一类计算量小到可以忽略但内存访问必须高效。1.2 训练视角和推理视角的错位很多人会困惑既然词嵌入是查表那为什么训练代码里还要定义权重矩阵、还有参与反向传播这是因为训练时每个token对应的行需要更新梯度本质是稀疏的矩阵更新——梯度只落在被采样到的那些行上。而在推理引擎里我们只需要前向查表完全没有梯度所以可以把嵌入层退化为一个纯粹的gather操作。这里有个实际工程含义推理引擎的嵌入层不应该保留任何训练用的接口设计。比如有些实现会把嵌入层当普通Linear层处理用矩阵乘去实现这在推理时纯属浪费。正确做法是直接读取权重buffer按行索引做拷贝或视图映射。1.3 嵌入向量的意义从语义离散到连续空间虽然这一章重点在工程实现但我还是想强调一下嵌入向量的语义角色。LLM之所以能理解语言就在于把离散的token映射到高维连续空间后语义相似的词在向量空间中距离更近。这个连续空间是后续所有Transformer层加工的基础。工程上一个有意思的现象嵌入层和输出层lm_head在很多模型里是权重共享的。也就是说模型末尾把hidden state映射回词表概率时用的权重矩阵和开头的嵌入矩阵是同一份数据。这个设计叫tied embedding我们在写推理引擎时加载权重文件时看到的token_embd.weight和output.weight可能对应同一块内存。处理不当的话权重重复加载会导致显存浪费甚至产生诡异的结果。2. 权重存储与内存布局嵌入层的一半工作量在这里2.1 形状、精度与磁盘加载的坑嵌入权重在模型文件里的形状一般是[vocab_size, hidden_size]也有反过来的版本[hidden_size, vocab_size]取决于训练框架。写推理引擎时必须先读取模型配置确认到底是哪种排列再决定加载时是否转置。这里有个实践建议嵌入层不需要像Transformer里的Linear层那样追求矩阵乘法的cache友好布局因为推理时它只按行读。如果文件里存的是[hidden_size, vocab_size]我们有两种选择一是加载时花一点时间转成[vocab_size, hidden_size]方便后续取行二是加载时不转取行时按列间隔访问。第二种在显存里其实很痛苦——每次取一行都要跳hidden_size个元素cache命中率极低。我最终选择了加载时统一转置。精度方面主流模型的嵌入层通常是fp16或bf16存储。推理引擎在CPU上一般用float做计算所以需要做精度转换。llama.cpp的做法是每个block定义自己的权重读取函数嵌入层因为是按行取转换损失可以忽略。我在CPU版本里直接用memcpy把原始值搬进临时缓冲再用半精度转单精度的函数批量处理。2.2 量化格式下的嵌入层被忽略的特例现成推理引擎里权重量化通常是重头戏4bit、8bit、各种group quantization。但很少有人提嵌入层在量化时的特殊处理。以GPTQ和AWQ为例这两类量化方案通常情况下不量化嵌入层和输出层或者只做极低比例的量化。为什么因为嵌入层直接决定token的初始表示量化误差会像“地基歪了”一样传导到每一层。有些引擎实现为了省显存硬要量化嵌入层结果就是模型困惑度明显上升生成质量肉眼可见地变差。如果你在做自己的推理引擎和量化方案集成时请务必保留嵌入层的完整精度。我踩过一次坑为了把整个模型压进显存把嵌入层也做了8bit量化跑出来回答就变得前言不搭后语。后来看了llama.cpp的源码才发现别人早就把这个排除在量化之外了。这事给我一个教训不要为了统一代码路径牺牲关键层的精度。2.3 token ID与向量维度的边界检查写嵌入层代码时最常见也最致命的问题是对token ID不做边界校验。正常人写代码时觉得“tokenizer输出的ID肯定在词表范围内”但真实推理中一旦遇到模型配置里没有覆盖的特殊token或者tokenizer版本和模型版本不匹配越界ID就会直接让引擎读到垃圾内存轻则输出乱码重则segfault。我在实现里加了一个开关在加载模型时记录vocab_size在取行前检查ID是否落在[0, vocab_size)内越界时统一映射到一个特殊token或者直接报错退出。这个检查在batch循环外层做一次就好不要在每一行取的时候做否则会拉慢速度。另一个细节是token ID的偏移。有些模型的词表里padding占位符不是0或者中间有一段未使用区间。如果直接用原始ID取行可能取到的是未训练的随机初始化行。遇到这种情况需要查配置文件里的bos_token_id、eos_token_id、pad_token_id以及是否存在特殊偏移。3. 手写实现从最简单的循环到高效Gather3.1 CPU端的直接实现先给一个最简单的CPU版本走的是“逐token取行拷贝精度转换”的路线。这个版本虽然不追求极致性能但逻辑清楚适合做基准。// 假设权重以fp16存储加载后统一转成fp32到weight_f32 void embedding_forward(float *output, const uint16_t *weight_half, const uint32_t *input_ids, const int seq_len, const int hidden_size, const int vocab_size) { for (int i 0; i seq_len; i) { uint32_t id input_ids[i]; if (id (uint32_t)vocab_size) { // 越界保护按实际场景处理 id 0; } const uint16_t *row weight_half (size_t)id * hidden_size; float *dst output (size_t)i * hidden_size; for (int j 0; j hidden_size; j) { dst[j] half_to_float(row[j]); } } }这个实现看起来人畜无害但性能上有两个明显缺点。第一half_to_float每个元素都做一次转换虽然现代CPU有向量化指令比如F16C但编译器不一定能自动识别这种逐元素的半精度转换并完美向量化。第二按行取的时候cache还算友好但循环内部没有利用hidden_size对齐的特点做批量处理。实际优化时我倾向于预先准备一个n对齐的临时缓冲然后对整行做向量化转换再memcpy到输出。在支持avx2的机器上批量转换的速度快很多。3.2 针对连续token ID的批量优化推理时有一个很常见的场景同一batch内不同序列的token ID各不相同但同一序列内输入部分prefill阶段的token ID往往连续密集。这带来一个优化空间如果某个batch行内有一大段ID是连续递增的我们可以直接按区间拷贝权重行而不是逐个ID单独访问。我在实现里加了一个简单的“连续段检测”遍历token ID时如果当前ID等于上一个ID加1就把这一段标记为连续区间对应权重行也可以连续读取。实测在长prompt的prefill阶段这个优化能省下不少取行开销。当然这个优化对显存版AI加速卡的帮助弱一些GPU上更依赖带宽和并行度但依然是合法的优化方向。3.3 GPU端的kernel怎么写到了GPU上嵌入层就是一个标准的gather kernel。核心思路是每个线程处理输出向量的一个分量根据token ID定位到权重行拷贝对应分量。__global__ void embedding_kernel(const float *weight, const int *input_ids, float *output, const int seq_len, const int hidden_size) { int idx blockIdx.x * blockDim.x threadIdx.x; int total seq_len * hidden_size; if (idx total) { int row idx / hidden_size; int col idx % hidden_size; int id input_ids[row]; output[idx] weight[id * hidden_size col]; } }这里有三个注意事项。第一idx / hidden_size和idx % hidden_size是整数除法速度较慢可以改成uint32_t的位运算技巧因为hidden_size基本都是2的幂或者至少对齐到多个32位。第二要避免每个线程重复读取id最好通过共享内存把当前block需要的token ID先载入。第三如果权重是fp16存储、输出需要fp32就涉及精度转换可以用__half2float内建函数。3.4 嵌入层之后与归一化和位置编码的衔接嵌入层的结果不会直接进入Transformer的自注意力层而是先经过一个LayerNorm不少模型用的是RMSNorm再加上位置信息或旋转位置编码RoPE。这里有个关键衔接问题嵌入向量和位置编码相加的顺序。像GPT-2这样的经典模型是先加绝对位置编码再进Transformer而Llama系列用的RoPE不改变输入向量而是在Q/K计算时通过旋转矩阵注入位置信息。所以嵌入层输出在进入第一个Transformer Block前往往先做一次RMSNorm。我的经验是这部分逻辑不要写在嵌入层里。嵌入层保持“纯取行输出向量”归一化和位置信息的处理交给下游的Block。这样模块职责单一加载不同模型时嵌套层的代码不需要改。4. 实测与踩坑记录我在这层栽过的跟头4.1 输出错位多了一个BOS token导致全文偏移有段时间模型跑出来的结果每个词都往后错一位例如输入“今天天气很好”输出变成“天天天气很好”这种重复或移位现象。排了很久才发现问题出在prefill阶段我在拼接prompt时手动加上了BOS token但加载模板时tokenizer已经在开头生成了一个BOS导致最终输入序列变成[BOS, BOS, token1, token2, ...]嵌入层输出的向量序列整体向右便宜了一个token位置。这类问题在推理引擎里特别隐蔽因为单个token错位不会报错但语义完全乱了。排查建议直接在嵌入层输出处打印出第一行向量对比模板或者对着token ID序列检查是否有重复BOS。4.2 共享权重下的显存重复加载其实上文提到过tied embedding这里分享一个真实案例我的第一个版本加载权重时按名字依次处理token_embd.weight和output.weight在磁盘上是两个条目但数据内容完全一样共享权重。我那时没做去重直接加载了两次2GB的嵌入层占用瞬间变成4GB对一个本就不宽裕的推理环境来说非常致命。后来我在加载器里加了一张表如果两个条目指向同一偏移区域就只加载一次并让后续引用共用同一块buffer。这个优化对显存好几十G的模型来说微不足道但对中小模型来说非常值。4.3 嵌入层性能占比与优化优先级很多人以为推理引擎的性能瓶颈一直在Transformer主体嵌入层不值一提。实测下来prefill阶段嵌入层的耗时占比大约是1%~3%生成阶段decode单token嵌入层占比可以忽略。所以优化时间不应该花在嵌入层的极致性能上而应该把重点放在它的功能和兼容性上。我的建议是嵌入层追求的是正确、稳定、清晰而不是快。代码应该好读、边界检查齐全、权重布局明了后面接Transformer时越省心越好。而且输出层lm_head常常和它共享权重理解嵌入层的存储方式对后面的logits计算也有直接帮助。4.4 特殊token与未训练token的处理策略真实使用中模型词表里有相当一部分token是“保留位置”比如未分配的编码、某些语言的特有符号。如果输入ID正好落在这些位置嵌入层会取到一个从未被训练过的随机向量导致模型输出明显变差。我的处理策略是加载模型配置时记录一个“有效token范围”在嵌入层对落在可疑区间的ID做告警或替换。这个开关默认关闭因为大多数场景不会触发但如果你在调试生成质量开着它有助于快速定位是不是tokenizer和模型版本不匹配。补充一个实操技巧在很多开源模型的tokenizer里特殊token放在词表末尾例如|endoftext|。如果使用者在输入文本里手动加了这些特殊符号嵌入层也能正常取行但后续注意力层可能因为pad的位置信息混乱导致效果异常这个要区分“嵌入层没问题”和“模型使用不对”。5. 和Transformer主体的衔接为下一章做准备嵌入层输出[seq_len, hidden_size]的浮点矩阵后接下来的流程对推理引擎的设计影响很大。不同模型的衔接方式有细微差别Llama系列嵌入层 → RMSNorm → 第一个Transformer Block。RoPE在Block内部的Q/K计算时应用。GPT-2嵌入层 → LayerNorm早期版本没有后期在第一个Block前加了pre-norm→ 加绝对位置编码 → Transformer Block。ChatGLM嵌入层直接进Transformer Block归一化在每个Block内部做。所以做推理引擎的时候不要把嵌入层和后续层紧紧耦合在一起。我最终实现的结构是tokenizer - input_ids - embedding - (pre_norm / rope_prep) - transformer blocks嵌入层永远是干净的“ID到向量”映射算子任何模型的特殊预处理都不应该侵入这层。这里还想多提一句RoPE与嵌入层的关系。RoPE不是把位置编码加到向量上而是在Q/K向量上按位置乘旋转矩阵。因为Q/K是从嵌入层输出经过多个线性变换得到的所以嵌入层本身不需要感知位置。这一点和绝对位置编码的实现差异很大如果混淆了写出来的引擎很难同时兼容两类模型。6. 最后分享一点这章之外的心得写推理引擎这个系列写到这里我最大的感受是真正决定一个引擎好不好维护的往往不是那些炫酷的注意力优化而是基础算子边界处理得干不干净。词嵌入这一层就是典型例子网上资料少看起来没难度但你一旦在共享权重、量化排除、边界检查上偷懒后续排错的成本会翻十倍。我给自己的编码习惯定了几条规矩放在这里供参考嵌入层代码单独一个文件不掺其他算子逻辑所有取行路径统一走同一个gather函数模型配置里的vocab_size、hidden_size、tie_word_embeddings必须显式读到结构体里不靠猜测每次跑模型前先用单token输入验证嵌入层输出是否和参考实现一致。下一章我们会进入Transformer Block本身从注意力机制开始讲。嵌入层作为这条链路的起点它的输出质量、内存布局、量化策略已经为后续所有计算打好了地基。地基稳不稳跑一遍模型就能感受到。