Prefill 与 Decode：LLM 推理的两个执行阶段-平芜编程栈

Prefill 阶段：一次性读完所有 Token

# Prefill 阶段的推理调用——AscendCL 同步执行importpyascaspaimportnumpyasnp pa.init()device=pa.set_device(0)model=pa.load_model("llama.om")# Prefill：一次性处理整个输入序列 [1, n] → [1, n, d]input_ids=np.array([[45,892,312,67,1289,34,567]],dtype=np.int64)input_tensor=pa.Tensor(input_ids,dtype=pa.int64)# 执行 Prefill——内部调用 GE 的 Prefill 图分支output_tensor=model.execute([input_tensor])# 同步：等所有 Block 跑完first_token=output_tensor.to_numpy()print(f"Prefill 完成，第一个 Token:{first_token.argmax(-1)[0]}")

Prefill 把[1, n]的输入序列一次性传给模型。40 个 Decoder Block 全部跑一次——每层算 Attention 时 Q、K、V 都是完整的[1, n, d]。计算量跟 n² 成正比。

Prefill 计算量 = 40 layers × (Q@K: 2×n×d² + Score@V: 2×n²×d + FFN: 4×n×d²) n=2048, d=4096 时：约 2.7 TFLOPs → 在 24 TFLOPS 的 NPU 上：约 110ms

CANN Runtime 在 Prefill 阶段用大 Tile 跑 Cube Unit——M 维度的 Batch=1 但乘以 n（序列长度），实际上 M=n，Cube 利用率可以做到 80%+。

KV Cache 的产生

// KV Cache 在 Prefill 阶段被创建和填充// 每层每个 Head 的 K 和 V 被缓存下来structKVCacheBlock{void*k_buffer;// [num_heads, seq_len, head_dim]void*v_buffer;// [num_heads, seq_len, head_dim]intcurrent_len;// 当前缓存的 Token 数};// Prefill 结束后 Cache 已满KVCacheBlock cache[40];// 40 layersfor(intlayer=0;layer<40;layer++){cache[layer].current_len=n;// 2048 个 Token 全部缓存在 Cache 中}

KV Cache 的内容：每层每个 Attention Head 的 K 矩阵和 V 矩阵。Prefill 结束后 Cache 已经保存了输入全部 Token 的 K 和 V。Decode 阶段只需要追加新 Token 的 K/V。

Decode 阶段：逐 Token 生成

# Decode 阶段——循环逐 Token 生成decoded_tokens=[first_token.argmax(-1)[0]]max_new_tokens=512forstepinrange(max_new_tokens):# 每次只输入上一个 Token，而非完整序列last_token=np.array([[decoded_tokens[-1]]],dtype=np.int64)token_tensor=pa.Tensor(last_token,dtype=pa.int64)# GE 自动切换到 Decode 图分支——单 Token 计算output=model.execute([token_tensor])next_token=output.to_numpy().argmax(-1)[0]decoded_tokens.append(next_token)# Runtime 在每步后自动更新 KV Cache——追加新 K/Vifnext_token==2:# EOS Tokenbreakprint(f"最终输出:{decoded_tokens}")# 解码 512 Token 约 640ms（1.25ms/Token）

Decode 每步只输入 1 个 Token。Attention 计算的不是Q @ K——K 已经在 Cache 里了。每步只算 Q，然后Q @ K_cache。计算量 = Prefill 的 1/n。

Decode 的 Cache 访问模式

// Decode 阶段的 Attention 计算——用 KV Cache 避免重复计算voiddecode_attention(float*Q,KVCacheBlock&cache,float*output,inthead,intstep){// Q 是当前 Token 的 Query：[1, head_dim]// K_cache 已经包含了之前所有 Token 的 Key：[step, head_dim]// 不需要重新算 K// Attention Score = Q @ K_cache^Tfloatscore[step];// score[i] = dot(Q, K_cache[i])for(inti=0;i<step;i++){score[i]=vector_dot(Q,cache.k_buffer+i*head_dim,head_dim);}// Softmax + @Vsoftmax(score,step);floatresult[head_dim]={0};for(inti=0;i<step;i++){vector_mac(result,cache.v_buffer+i*head_dim,score[i],head_dim);}// 更新 Cache——追加当前 Token 的 K 和 Vmemcpy(cache.k_buffer+step*head_dim,current_k,head_dim*sizeof(float));memcpy(cache.v_buffer+step*head_dim,current_v,head_dim*sizeof(float));cache.current_len=step+1;memcpy(output,result,head_dim*sizeof(float));}

Cache 的 key 是"空间换时间"：不用重新算历史的 K 和 V——但 Cache 本身占用显存。LLaMA-13B 在 n=4096 时 Cache 占 3.2GB，Batch=8 时 25.6GB——超过模型参数本身（26GB）。

Runtime 如何调度两个阶段

// GE 在模型加载时编译了两套执行计划——Prefill 和 Decode// Runtime 根据输入 Shape 自动切换voidruntime_dispatch(Model*model,Tensor*input,Tensor*output){intbatch=input->shape[0];intseq_len=input->shape[1];if(seq_len>1){// 输入有多个 Token → Prefill 模式// 执行 Prefill 图分支——全量 Attentionge_execute(model->prefill_graph,input,output);// Prefill 结束后 Cache 已满}else{// 输入只有 1 个 Token → Decode 模式// 执行 Decode 图分支——Cache 读取 + 单 Token Attentionge_execute(model->decode_graph,input,output);// 每步更新 Cachekv_cache_append(model->cache,input->k,input->v);}}

GE 的图编译在两个路径上的差异：

Prefill 图：FlashAttention 用块状实现——Score 矩阵整块计算
Decode 图：FlashAttention 用 Paged 实现——逐 Block 读取 KV Cache

性能对比

# Prefill vs Decode 的实测延迟prefill_stats={"n=128":12,# ms"n=512":38,# ms"n=2048":110,# ms"n=4096":205,# ms}# Prefill 延迟随 n 近线性增长decode_stats={"step=1":1.2,# ms/Token"step=100":1.3,"step=500":1.5,"step=2000":2.1,}# Decode 延迟随 step 轻微增长——KV Cache 变长后 Attention 的搬运量增加