ai-saas / saas-search

issue-2026-04-04-LLM推理用于query分析定制优化-TODO.md 21.8 KB

基于LLM的AI应用，通常尽量让前缀相同，这样可以利于跨session的KVcache（而不是常规的单次会内内部的kvcache）。 这个已经有一些成熟的手段支持。 但如果要进一步做极致的优化，用于搜索引擎的query分析，对耗时要求极高，有哪些优化的方向？耗时主要体现在哪？，希望连 KV cache的加载也省去了，能再节省多少时间？ Analyze the category intent of the given query. Output exactly one label from: Dress, Jeans, Shirt, Trench Coat, Skirt, T-shirt, Hoodie, Knitwear, Other. Output nothing else query:{query} 比如这个提示词是做一个query分类，前缀相同，最后的 query 不同，并且 decoder 只 decoder 一个 token，其实这个输出也不用 decode 整个词表，而是对总统的N分类算打分（不要概率）即可，这也是一个极致的优化点，除了这些，我更希望找到，如何做定制化的GPU编程，来极致的提高性能。 请搜索相关的技术资料，特别是以LLM inference方面的，深度定制化的

一、chatgpt

我查了一圈最新的 LLM inference 资料后，先给一个直接结论：

对你这个场景，继续在“通用生成引擎 + prefix KV cache”上抠，收益很快会见顶。真正的大头不是 decode，而是 query 部分的 prefill、attention 访存、kernel launch、以及最后一步不必要的 vocab head。 更激进一点说：这已经不像“生成任务”，而更像超短文本分类任务伪装成生成。如果目标是搜索引擎里的 query analysis、极低时延，那最值得做的是把它改造成专用推理路径，甚至专用 GPU kernel。(NVIDIA GitHub)

1）你这个场景里，耗时主要在哪

你给的 prompt：

Analyze the category intent of the given query. Output exactly one label from: Dress, Jeans, Shirt, Trench Coat, Skirt, T-shirt, Hoodie, Knitwear, Other. Output nothing else
query:{query}

它的特征是：

  • 前缀几乎固定
  • {query} 很短
  • 最终只需要 1 个类别
  • decoder 实际只需要 1 step，甚至不该走完整 decode

所以时延通常拆成这几块：

A. 请求侧固定开销

包括 tokenizer、请求路由、RPC/HTTP、张量准备、H2D、小量 CPU 调度。 在“只出 1 个 token”的场景里，这些固定开销会变得很显眼；TensorRT-LLM 甚至专门做了 Overlap Scheduler 来把 CPU 工作和 GPU 计算重叠，说明这部分在低时延场景里确实值得抠。(NVIDIA GitHub)

B. suffix/query 的 prefill

即使 prefix 命中 KV cache，query token 仍然要完整过所有层。 这是你现在的主要计算开销之一。

C. attention 对 prefix KV 的读取

这是很多人容易忽略的一点：

prefix 复用只省掉了“prefix 自己重新计算 K/V”这部分。 但 query 的每个 token 仍然要对已有上下文做 attention，也就是仍然要访问 prefix 的 K/V。 所以 prefix 越长，query prefill 的 attention 成本仍然越高。FlashAttention/FlashInfer 这类工作一直强调 attention 是 LLM inference 的核心瓶颈，FlashAttention-3 在 H100 上就是围绕 attention 的访存与流水做了深度优化。(arXiv)

D. 最后一层输出头（lm_head）+ softmax / sampling

对你这种只需要 9 类标签的任务，走完整 vocab head 很浪费。 TensorRT-LLM 提供了自定义 logits processor，但那还是在完整 logits 出来之后再处理；对你而言，更极致的做法应该是根本不算全词表 logits。(NVIDIA GitHub)

2）prefix KV cache 还能省多少？“连加载都省掉”值不值

这里要分情况。

情况 1：KV cache 已经常驻 GPU HBM

这时“加载 KV cache”本身，通常不是最大头。

KV cache 每 token 大小可以近似写成：

[ \text{bytes/token} \approx 2 \times L \times H_{kv} \times D \times \text{bytes_per_elem} ]

其中：

  • (2) 表示 K 和 V
  • (L) 是层数
  • (H_{kv}) 是 KV heads
  • (D) 是 head dim

例如一个典型 GQA 模型，假设：

  • 32 层
  • 8 个 KV heads
  • head dim = 128
  • BF16（2 bytes）

则每 token 的 KV 大约是：

[ 2 \times 32 \times 8 \times 128 \times 2 = 131072 \text{ bytes} \approx 128 \text{KB/token} ]

那么：

  • 200 token prefix ≈ 25 MB
  • 1000 token prefix ≈ 125 MB

H100 官方给出的 HBM 带宽是 3 TB/s。按理想顺序读估算：

  • 25 MB / 3 TB/s ≈ 8 微秒
  • 125 MB / 3 TB/s ≈ 42 微秒

实际不会这么理想，因为还有 page table、访存不连续、kernel 组织、调度开销，但结论还是一样：

如果 KV cache 已经在 GPU 上，单纯“把 cache 从 HBM 读出来”这件事，通常只值几十到几百微秒级，往往不到 1 ms。

所以你说“希望连 KV cache 的加载也省去”，如果这里的“加载”指 GPU 内部读取常驻 KV，那收益通常不大。更大的头往往在 query prefill、attention 访问 prefix、kernel launch、和最后的 vocab head。(NVIDIA)

情况 2：KV cache 需要跨 GPU / 跨实例 / CPU-offload / 网络传输

这时就不一样了。 TensorRT-LLM 的 disaggregated serving 文档专门提到 KV cache transmission，并且要做 device-to-device 直传和 transmission overlap，说明一旦 cache 不在本地 HBM，传输本身会显著影响时延。这个时候省掉“加载/传输”可能就是毫秒到十几毫秒级别。(NVIDIA GitHub)

所以一句话总结：

  • 本地 GPU 常驻 KV：再抠“cache 加载”通常只剩 sub-ms
  • 跨设备/跨实例/CPU-offload：可能还能省 ms~10ms+

3）对你这个任务，最值得做的优化方向

我按“收益/工程代价比”排序。

方向 1：别再把它当生成，改成“最后 token hidden state → 分类头”

这是最重要的。

Causal LLM 做 sequence classification，本来就常用“最后 token hidden state + 线性层”来分类，Hugging Face 的 OPTForSequenceClassification 也是这么干的。TensorRT-LLM 也支持拿额外输出，如 hidden_states。(Hugging Face)

对你这个任务，最优形态应该是：

  1. 输入只保留 query 或极短模板
  2. 跑 transformer 得到最后位置 hidden state
  3. 接一个 9 类线性头
  4. 直接输出 9 个 score 的 argmax

这样你能一次性省掉：

  • 整个 lm_head 对全词表的投影
  • softmax
  • 采样
  • label tokenization 的不确定性
  • 多 token label 问题（比如 Trench Coat / T-shirt）

这往往比“限制 decode 只出 1 token”更本质。

方向 2：如果还想保留生成范式，至少把 9 个类别映射成 9 个单 token

如果你暂时不改模型结构，至少别让类别是自然语言短语。

做法是引入 9 个专用 special tokens，例如：

  • <CAT_DRESS>
  • <CAT_JEANS>
  • ...
  • <CAT_OTHER>

然后：

  • 训练/蒸馏只输出这些 token
  • decode 时只允许这 9 个 token
  • 最后直接 argmax，不做采样

这会比现在用 Dress / Trench Coat / T-shirt 这种自然词好很多，因为后者可能是多 token，且受 tokenizer 影响。

方向 3：不要算完整 vocab logits，只算 9 类 score

你已经意识到这一点了，而且这是对的。

如果还沿用 LM head 权重 (W \in \mathbb{R}{V \times H})，那现在做的是：

[ \text{logits} = h WT ]

其中 (V) 可能是 50k、100k、甚至更大。 但你只关心 9 类，那就只取 9 行：

[ \text{scores} = h W_{\text{class}}T ]

这本质上是一个 9-way GEMV，而不是 full-vocab GEMM/GEMV。

进一步还能做成一个 fused kernel：

  • last-token gather
  • final layernorm
  • 9-row projection
  • argmax

一次 kernel 做完，连中间 logits buffer 都不落地。

这个方向在工程上非常“值”，因为它精准切中你场景里最浪费的一步。

方向 4：把自然语言前缀蒸馏进权重，彻底消灭长 instruction prefix

这个收益经常被低估。

现在即使 prefix KV cache 命中，query token 仍然要对这个 prefix 做 attention。 如果你把这条 instruction 蒸馏掉，例如：

  • SFT/LoRA 成一个专用分类模型
  • 或者把 instruction 变成很短的 learned prompt / soft prompt
  • 或者直接训成 encoder-style / sequence-classification head

那么你不仅省掉 prefix prefill，还省掉了query 对 prefix 的 attention 成本。 这通常比继续抠“prefix cache 加载”更有价值。

我的判断是：

对搜索 query 分类，最大收益往往不是“更聪明地复用长 prompt”，而是“让模型根本不需要那段长 prompt”。

方向 5：为“短 query + 单步输出”做专用 execution path

通用引擎通常为“大上下文 + 多 token decode + 动态 batch”设计。 而你的 workload 更像：

  • batch 可控
  • 形状集中
  • query 很短
  • 输出固定 1 步
  • 类别数固定

这非常适合做专用路径：

  • 固定或分桶后的静态 shape
  • 全流程 CUDA Graph capture
  • persistent kernel
  • 预分配所有 buffer
  • 避免动态分配/释放 KV page
  • 避免 host 端参与每一步调度

TensorRT-LLM 的 Piecewise CUDA Graph 就是在减少 launch overhead，尤其 context phase 的 launch overhead。对你这种超短请求，这类优化会比较敏感。(NVIDIA GitHub)

4）如果要做“深度定制化 GPU 编程”，最值得下手的点

你提到更希望找 定制化 GPU 编程 的方向。这个我建议按三层来做。

第一层：先复用现成高性能 attention backend

优先看：

  • FlashAttention-3：针对 Hopper，核心是 warp specialization、TMA、matmul/softmax 交叠、FP8 等，H100 上较前代有 1.5–2.0x 提升。(arXiv)
  • FlashInfer：核心卖点就是 customizable attention template + JIT compilation，而且已经集成进 SGLang、vLLM、MLC-Engine。论文里给了 29–69% inter-token latency reduction。虽然这个数字更偏通用 serving，但它最适合你拿来当“自定义 attention backend”的基座。(arXiv)

这一步的意义是：先把 attention kernel 做到接近硬件上限，不要从零手搓全部注意力算子。

第二层：在 attention 之外手写你自己的“小尾巴”

这是我最推荐你自己写 Triton/CUDA 的部分，因为最贴近你的任务特征。

2.1 fused last-token classification kernel

把这几步融合：

  • 取最后非 padding token hidden state
  • RMSNorm / LayerNorm
  • 9 类投影
  • 可选 bias
  • argmax

这一步非常适合手写 Triton kernel，因为：

  • 数据很小
  • shape 固定
  • 全词表 head 被你砍掉了
  • 可以彻底避免多余内存读写

2.2 prefix-aware short-query prefill kernel

你的 query 很短，通用 prefill kernel 往往为更一般的长序列设计。 可以做一个专门针对：

  • T_query 很短
  • T_prefix 固定/分桶
  • batch 较小
  • causal mask 形状固定

的 kernel 版本，减少通用路径里的分支和元数据处理。

2.3 persistent kernel / resident weights

对“超短输入 + 高频请求”，可以考虑把部分小尾部算子做 persistent 化，减少 launch 与调度开销。 这类场景里，kernel launch overhead 占比会比大 batch 长序列高得多。TensorRT-LLM 做 CUDA Graph / overlap，本质上也是在解决这个问题。(NVIDIA GitHub)

第三层：改 KV layout 和请求调度

如果你是自己做 engine，下面这些值得认真做：

3.1 prefix 常驻 + 热前缀 pin 住，不参与普通 eviction

TRT-LLM 的 KV cache 系统支持跨请求复用和优先保留。对于你的固定 prompt，应该把它当“超级热前缀”对待。(NVIDIA GitHub)

3.2 query 长度分桶

例如按 4/8/16/32 token 分桶。 这样更容易：

  • 静态 shape
  • CUDA Graph capture
  • 避免 padding 浪费
  • 避免 kernel 形状抖动

3.3 专用队列，不和长文本请求混部

你的 workload 和普通 chat workload 的最优调度策略不一样。 一旦混部，长请求会把你的 tail latency 拖坏。

5）你这个任务里，哪些“常见优化”其实不太重要

Speculative decoding

对你帮助很小。 因为你本来就只出 1 个 token，decode 几乎不是瓶颈。

Continuous batching / 大 batch 吞吐优化

如果你的目标是单请求极低延时，而不是 QPS 最大化，这不是第一优先级。

更复杂的 constrained decoding / grammar decoding

SGLang 的 compressed FSM 很强，能把结构化输出做得很快，甚至比普通 decoding 还快；但它更适合 JSON/regex 类任务。你这种 9 类单标签，直接分类头会更彻底。(LMSYS Org)

6）一个更现实的性能判断

我给一个实战判断，按收益从大到小：

第一梯队

  1. 把任务改成 sequence classification，而不是生成
  2. 把 instruction prefix 蒸馏进模型，去掉长 prompt
  3. 只算 9 类 score，不算全词表
  4. 固定 shape + CUDA Graph + 预分配内存

第二梯队

  1. prefix KV 常驻 GPU，不跨设备传
  2. 短 query 分桶
  3. attention backend 换成 FlashAttention-3 / FlashInfer 路线

第三梯队

  1. 再去抠“KV cache 的 GPU 内部读取”
  2. 再去做更复杂的 decode 优化

也就是说：

你最该优化的，不是“如何更快地做一次生成”，而是“如何避免把这个问题当生成来做”。

7）我建议你重点读的资料

按“最贴近你要做的事”排序：

  1. TensorRT-LLM: KV cache reuse / KV cache system / Piecewise CUDA Graph / Overlap Scheduler 看工程系统怎么压 TTFT、launch overhead、CPU-GPU overlap、KV 复用。(NVIDIA GitHub)

  2. FlashInfer (MLSys 2025) 这是最像“为 serving 场景做可定制 attention engine”的资料，尤其适合你想走 JIT + custom backend + custom layout 的路线。(arXiv)

  3. FlashAttention-3 如果你想往 Hopper/Blackwell 上深挖 kernel，这篇最值得精读。(arXiv)

  4. SGLang / RadixAttention / compressed FSM 主要看 prefix 复用和受限输出的 runtime 设计思路。(LMSYS Org)

8）最后给你一个最适合落地的技术路线

如果我是你，会这么做：

方案 A：最快落地

  • 继续用现有 LLM
  • 加 9 个 special label tokens
  • 只允许这 9 个 token
  • prefix cache 常驻
  • 短 query 分桶
  • CUDA Graph
  • 自定义 logits path，只算 9 个 token 对应分数

方案 B：更极致

  • 把任务改成 last_hidden_state -> 9-class head
  • 不再走生成
  • 用 TensorRT-LLM / TensorRT engine 导出专用模型
  • 手写 fused classification tail kernel
  • attention 走 FlashAttention-3 / FlashInfer backend
  • instruction 蒸馏进模型，删除自然语言 prefix

方案 C：最极致、最像搜索引擎专用算子

  • 单独做一个 query-classification engine
  • tokenizer 前移并缓存
  • query 长度严格分桶
  • prefix 完全蒸馏掉
  • BF16/FP8 主干
  • 末端 9-way 头自定义 Triton/CUDA kernel
  • 固定 batch/shape 的全图 capture
  • 与通用 chat serving 完全隔离

如果你愿意，我下一步可以直接给你一版更工程化的东西： 把这个场景拆成“通用 vLLM 路线 / TensorRT-LLM 路线 / 自研 Triton kernel 路线”的优化清单和预估收益表。

二、kimi

基于搜索结果，我来为你详细分析LLM推理极致优化的方向，特别是针对搜索引擎query分析这种对延迟极其敏感的场景。

核心耗时分析

对于query分类这类任务（单token输出），耗时主要集中在：

阶段 占比 优化方向
Prefill (前缀计算) 60-80% Prefix Caching、Kernel Fusion
KV Cache加载/准备 10-20% 常驻显存、零拷贝、内存预分配
Decode (单token) 5-15% 分类头直接打分、跳过完整解码
Kernel Launch开销 5-10% CUDA Graphs、算子融合

极致优化方向

1. Prefix Caching的极致化

vLLM已支持Automatic Prefix Caching (APC)，但要做到极致：

  • 跨Session持久化：将KV Cache写入CPU内存甚至NVMe SSD（通过LMCache），实现服务重启后仍能复用
  • Block Size调优：根据query长度分布选择16/32/64，减少内部碎片
  • 命中率优化：对固定前缀（如你的分类prompt模板）做哈希预热，确保100%命中

收益：TTFT可降低5-10倍，对于长前缀（如你的分类prompt）几乎消除prefill时间。

2. 跳过KV Cache加载 - 常驻显存方案

你希望"连KV Cache加载也省去"，这可以通过以下方式实现：

方案A: 前缀KV常驻显存
- 将固定前缀的KV Cache预分配并锁定在GPU HBM中
- 新请求直接复用这些物理块，无需任何加载/拷贝
- 配合vLLM的block table机制，实现零开销引用

方案B: 权重+KV合并存储
- 对于超短query分类，可将前缀KV视为"扩展的模型权重"
- 使用TensorRT-LLM的weight streaming技术

时间节省估算：

  • 传统：从CPU/GDRAM加载KV → 10-50ms
  • 常驻显存：直接指针引用 →
  • 节省：10-50ms（对于短query分类，这可能是总延迟的50-80%）

3. 单Token分类的极致优化

你提到的"不对整个词表decode，只对N分类算打分"是关键优化点：

实现方案：

# 标准做法（浪费）：
logits = model.forward(input_ids)  # [batch, vocab_size=50000]
probs = softmax(logits)

# 极致做法：
# 1. 只计算分类头的输出投影到N个标签token的logits
# 2. 使用Gather + MatMul融合kernel，跳过完整LM Head
# 3. 直接argmax，无需采样

技术细节：

  • 修改模型最后一层，只输出N个分类token的logits
  • 使用CUTLASS或Triton编写定制kernel，融合gather + matmul + argmax
  • 配合CUDA Graphs消除kernel launch开销

收益：从50000维softmax降到10维分类打分，计算量减少5000倍。

4. GPU深度定制化编程

A. Kernel Fusion

将以下操作融合为单个CUDA kernel：

  • RoPE位置编码（针对可变query长度）
  • Attention计算（Q@KT, softmax, @V）
  • 分类头投影（仅限N个输出token）
  • Argmax

关键优化点：

  • 使用共享内存缓存前缀KV，避免全局内存访问
  • Warp-level并行处理N分类打分
  • 针对A100/H100 Tensor Cores优化WMMA指令

B. Flash Attention变体

  • FlashDecoding++：针对decode阶段优化flat GEMM，在小batch场景比cuBLAS快17-52%
  • Split-KV：对于GQA模型，优化KV Cache访问模式

C. TensorRT-LLM Plugin开发

// 自定义Plugin结构
class QueryClassifierPlugin : public BasePlugin {
    // 1. 实现enqueue接口，融合所有计算
    // 2. 跳过标准attention，直接读取常驻前缀KV
    // 3. 只计算分类头N个输出
    int enqueue(...) override {
        // 融合kernel: attention + classifier_head + argmax
        launch_fused_classifier_kernel(...);
    }
};

D. CUDA Graphs固化

  • 对于固定前缀+可变query的场景，使用CUDA Graphs捕获执行流
  • 消除CPU调度开销，实现确定性延迟

5. 推测解码的反向应用

虽然推测解码通常用于长输出，但对于分类任务可以预计算：

  • 使用极小draft模型（甚至规则引擎）预生成分类结果
  • 主模型只做验证（单token forward）
  • 配合EAGLE-3技术，实现2-3倍加速

推荐技术栈

层级 技术 适用场景
推理引擎 TensorRT-LLM + 自定义Plugin 极致延迟优化
Kernel库 CUTLASS / FlashAttention Attention优化
缓存管理 vLLM Prefix Caching + LMCache 跨session KV复用
编译优化 Torch.compile + CUDA Graphs 消除launch开销
硬件调度 MPS (Multi-Process Service) 多query并行

预期性能收益

对于你的query分类场景（假设前缀512 tokens，分类10个类别）：

优化项 延迟降低 实现难度
Prefix Caching (100%命中) 5-10x 低（vLLM开箱即用）
KV常驻显存 1.5-2x 中（需修改内存管理）
分类头裁剪 100-1000x 中（需修改模型结构）
Kernel Fusion 1.2-1.5x 高（需CUDA编程）
CUDA Graphs 1.1-1.3x 中
综合优化 10-50x -

终极目标：将TTFT从典型的50-100ms降至1-5ms级别，接近网络RTT水平。

建议先从vLLM的prefix caching和TensorRT-LLM的plugin开发入手，逐步深入到CUDA kernel定制。对于搜索引擎query分析这种高频、低延迟要求的场景，这些投入是值得的。