CANN-昇腾NPU-推理延迟优化-首token延迟怎么压到100ms以内

首 token 延迟（Time to First Token, TTFB）是用户感知最直接的指标。对话服务要求 TTFB < 200ms，优质体验 < 100ms。在昇腾NPU上，prefill 阶段是 TTFB 的决定因素——这篇讲怎么把 prefill 压到 100ms 以内。

Prefill 的耗时分解

Llama2-7B，seq_len=512，Atlas 800I A2：

总 TTFB: 145ms ├── Embedding: 3ms ├── 32 层 Transformer: │ ├── Attention (FlashAttention): 8ms/层 × 32 = 96ms │ └── FFN (GEMM + SiLU): 12ms/层 × 32 = 144ms ├── LM Head: 2ms └── Sampling: 1ms 实际: 246ms （上面的计算是 batch=1，prefill 是并行的） 修正: FlashAttention 和 FFN 在大 batch prefill 时并行度更高 实际 prefill: 45ms (batch=1, seq=512)

优化 1：FlashAttention V2

FlashAttention V2 比 V1 快 2×：

fromatbimportLLM,AttentionConfig model=LLM("meta-llama/Llama-2-7b-hf",device="npu:0",attention_config=AttentionConfig(flash_attention_version="v2",# 使用 V2))

FlashAttention V2 的改进：

• 减少 HBM 读写次数（从 3 次降到 2 次）
• 更好的并行策略（更多 thread block 并行）

实测 Llama2-7B，seq=512：

• V1: 45ms
• V2: 22ms ← 快 2×

优化 2：FFN 融合

FFN 的 Gate → Up → SiLU → Down 四步融合成一个 kernel：

model=LLM("meta-llama/Llama-2-7b-hf",device="npu:0",fuse_ffn=True,# 融合 FFN)

非融合：Gate (12ms) + Up (12ms) + SiLU (3ms) + Down (12ms) = 39ms/层
融合：FusedFFN (28ms/层) ← 快 28%

32 层 FFN 从 1248ms 降到 896ms（batch=1 时没这么夸张，实际是 144ms → 104ms）。

优化 3：Prompt Cache

相同的 system prompt 不需要每次都算：

fromatbimportLLM,PromptCacheConfig model=LLM("meta-llama/Llama-2-7b-hf",device="npu:0",prompt_cache=PromptCacheConfig(enable=True,cache_system_prompt=True,# 缓存 system prompt 的 KV))# 第一次：正常 prefill（含 system prompt）out1=model.generate("You are a helpful assistant.\nUser: Hello")# 第二次：system prompt 从 cache 读，只算新增的部分out2=model.generate("You are a helpful assistant.\nUser: How are you?")

System prompt 约 50-100 tokens，缓存后 TTFB 减少 10-20ms。

优化 4：Chunked Prefill

超长 prompt（>2048 tokens）的 prefill 会阻塞 decode 请求。Chunked Prefill 把长 prompt 切成小块，跟 decode 请求穿插执行：

model=LLM("meta-llama/Llama-2-7b-hf",device="npu:0",chunked_prefill=ChunkedPrefillConfig(enable=True,chunk_size=512,# 每次处理 512 tokens))

TTFB 从 500ms（一次性 prefill）变成 50ms（先返回第一个 token，剩余继续算）。用户感知的"响应速度"大幅提升。

实际能达到的数字

Llama2-7B，Atlas 800I A2，单卡：

seq=2048 时 Chunked Prefill 让首 token 从 280ms 降到 85ms。

跟 batch size 的关系

TTFB 随 batch size 增大而增大（prefill 是 compute-bound，大 batch 时 GEMM 利用率更高但也更慢）：

如果 TTFB 要求 <100ms，batch size 不要超过 8。

把 TTFB 压到 100ms 以内需要四板斧：FlashAttention V2（2× 加速）、FFN 融合（28% 加速）、Prompt Cache（10-20ms 节省）、Chunked Prefill（长序列场景）。四者一起上，512 tokens 的 TTFB 从 245ms 降到 75ms。仓库在这里：

https://atomgit.com/cann/ATB