docs/issues/issue-2026-04-04-LLM%E6%8E%A8%E7%90%86%E7%94%A8%E4%BA%8Equery%E5%88%86%E6%9E%90%E5%AE%9A%E5%88%B6%E4%BC%98%E5%8C%96-TODO.md

基于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])
---
## 1）你这个场景里，耗时主要在哪
你给的 prompt：
```text
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][2])
### 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][3])
### D. 最后一层输出头（lm_head）+ softmax / sampling
对你这种只需要 9 类标签的任务，走完整 vocab head 很浪费。
TensorRT-LLM 提供了自定义 logits processor，但那还是在**完整 logits 出来之后再处理**；对你而言，更极致的做法应该是**根本不算全词表 logits**。([NVIDIA GitHub][4])
---
## 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][5])
### 情况 2：KV cache 需要跨 GPU / 跨实例 / CPU-offload / 网络传输
这时就不一样了。
TensorRT-LLM 的 disaggregated serving 文档专门提到 KV cache transmission，并且要做 device-to-device 直传和 transmission overlap，说明**一旦 cache 不在本地 HBM**，传输本身会显著影响时延。这个时候省掉“加载/传输”可能就是**毫秒到十几毫秒**级别。([NVIDIA GitHub][6])
所以一句话总结：
* **本地 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][7])
对你这个任务，最优形态应该是：
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 W^T
]
其中 (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][8])
---
## 4）如果要做“深度定制化 GPU 编程”，最值得下手的点
你提到更希望找 **定制化 GPU 编程** 的方向。这个我建议按三层来做。
### 第一层：先复用现成高性能 attention backend
优先看：
* **FlashAttention-3**：针对 Hopper，核心是 warp specialization、TMA、matmul/softmax 交叠、FP8 等，H100 上较前代有 **1.5–2.0x** 提升。([arXiv][3])
* **FlashInfer**：核心卖点就是 **customizable attention template + JIT compilation**，而且已经集成进 SGLang、vLLM、MLC-Engine。论文里给了 **29–69% inter-token latency reduction**。虽然这个数字更偏通用 serving，但它最适合你拿来当“自定义 attention backend”的基座。([arXiv][9])
这一步的意义是：先把 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][8])
---
### 第三层：改 KV layout 和请求调度
如果你是自己做 engine，下面这些值得认真做：
#### 3.1 prefix 常驻 + 热前缀 pin 住，不参与普通 eviction
TRT-LLM 的 KV cache 系统支持跨请求复用和优先保留。对于你的固定 prompt，应该把它当“超级热前缀”对待。([NVIDIA GitHub][10])
#### 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][11])
---
## 6）一个更现实的性能判断
我给一个实战判断，按收益从大到小：
### 第一梯队
1. **把任务改成 sequence classification，而不是生成**
2. **把 instruction prefix 蒸馏进模型，去掉长 prompt**
3. **只算 9 类 score，不算全词表**
4. **固定 shape + CUDA Graph + 预分配内存**
### 第二梯队
5. **prefix KV 常驻 GPU，不跨设备传**
6. **短 query 分桶**
7. **attention backend 换成 FlashAttention-3 / FlashInfer 路线**
### 第三梯队
8. **再去抠“KV cache 的 GPU 内部读取”**
9. **再去做更复杂的 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][1])
2. **FlashInfer (MLSys 2025)**
   这是最像“为 serving 场景做可定制 attention engine”的资料，尤其适合你想走 JIT + custom backend + custom layout 的路线。([arXiv][9])
3. **FlashAttention-3**
   如果你想往 Hopper/Blackwell 上深挖 kernel，这篇最值得精读。([arXiv][3])
4. **SGLang / RadixAttention / compressed FSM**
   主要看 prefix 复用和受限输出的 runtime 设计思路。([LMSYS Org][12])
---
## 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 路线”的优化清单和预估收益表。**
[1]: https://nvidia.github.io/TensorRT-LLM/advanced/kv-cache-reuse.html "KV cache reuse — TensorRT-LLM"
[2]: https://nvidia.github.io/TensorRT-LLM/1.1.0rc4/features/overlap-scheduler.html "Overlap Scheduler — TensorRT-LLM"
[3]: https://arxiv.org/abs/2407.08608 "[2407.08608] FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision"
[4]: https://nvidia.github.io/TensorRT-LLM/latest/examples/llm_logits_processor.html "Control generated text using logits processor — TensorRT LLM"
[5]: https://www.nvidia.com/en-us/data-center/h100/ "H100 GPU | NVIDIA"
[6]: https://nvidia.github.io/TensorRT-LLM/features/disagg-serving.html "Disaggregated Serving — TensorRT LLM"
[7]: https://huggingface.co/docs/transformers/model_doc/opt "OPT · Hugging Face"
[8]: https://nvidia.github.io/TensorRT-LLM/latest/features/torch_compile_and_piecewise_cuda_graph.html "Torch Compile & Piecewise CUDA Graph — TensorRT LLM"
[9]: https://arxiv.org/abs/2501.01005 "[2501.01005] FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving"
[10]: https://nvidia.github.io/TensorRT-LLM/latest/features/kvcache.html?utm_source=chatgpt.com "KV Cache System — TensorRT LLM - nvidia.github.io"
[11]: https://www.lmsys.org/blog/2024-02-05-compressed-fsm/ "Fast JSON Decoding for Local LLMs with Compressed Finite State Machine - LMSYS Blog | LMSYS Org"
[12]: https://www.lmsys.org/blog/2024-01-17-sglang/ "Fast and Expressive LLM Inference with RadixAttention and SGLang - LMSYS Blog | LMSYS Org"
# 二、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
- 常驻显存：直接指针引用 → <0.1ms
- **节省：10-50ms（对于短query分类，这可能是总延迟的50-80%）**
### 3. **单Token分类的极致优化**
你提到的"不对整个词表decode，只对N分类算打分"是关键优化点：
**实现方案**：
```python
# 标准做法（浪费）：
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@K^T, 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开发**
```cpp
// 自定义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分析这种高频、低延迟要求的场景，这些投入是值得的。