[Beam Search Optimization] Memory-Efficient Beam Search for multi-head attention
Author: brian030128
Pull request: no
Labels: name: Feature request

name: Generation

### Feature request

Adding a [Trie-Based Beam Search](https://www.arxiv.org/abs/2502.00085) implementation.


### Motivation

The current beam search implementation in Hugging Face's transformers library maintains separate KV caches for each beam candidate, even when they share common prefixes. For example, with a beam width of 4, identical prompt are redundantly stored 4 times, leading to excessive GPU memory usage and suboptimal inference speed due to memory bandwidth constraints.
The memory benchmark can be found in the paper above, where trie-based decoding method uses way less memory and slightly improves speed. 

This enhancement requires no additional libraries or hardware dependencies - it's purely an algorithmic improvement that can be implemented as a drop-in replacement while preserving the existing API. The optimization aligns with the growing industry focus on efficient inference and could benefit the many production systems currently using beam search for higher quality text generation.

### Your contribution

If this looks good for the maintainers, we're willing to implement and contribute this feature. We would appreciate feedback on the integration approach and are happy to provide additional benchmarks. We're ready to contribute both implementation and documentation.

Comments:
cc @gante . Btw, I am not sure if anyone already discussed the automatic prompt caching as a feature. I know we have been working to make something similar possible per user request, by cloning the cache. Just fyi since it is not same and kinda relevant to the linked paper

@zucchini-nlp  Hi, any updates or thoughts on this?

AFAIK prompt caching might be implemented as a general feature in transformers, some time in the future. I am not 100% sure about what are the plans, so gentle ping @gante again :)

Hi, thanks for the reply. This is not prompt caching. Instead, it is another decoding strategy for beam search in multi-head attention models.

Hi @brian030128 

Yeah, depending on the implementation details this would be desirable (and should be the default 👀). 

Adding a full decoding strategy would not be desirable -- it has a very high long-term maintenance cost for our team, and I'm currently designing an alternative solution to ease the inclusion of new generation techniques, hosting code on the HF Hub. After this hosting solution is complete, you can host your decoding method on the HF Hub and we should be able to run it like `model.generate(**model_inputs, trust_remote_code=True, recipe="brian030128/beam_method")`.

Nevertheless, and bear in mind I have only briefly scanned the paper, if it can be coded as a `Cache` class, I would welcome it in `transformers` 🤗 At `.update()` time it would convert the squashed cache into the full cache for the layer. It may need a new method to squash the cache, called inside `_beam_search`, but I'd be happy to include this pattern exception.

A final note: we're finalizing a refactor of beam search [here](https://github.com/huggingface/transformers/pull/35802).

Hi everyone,

I was going through the Trie-Based Beam Search paper and it looks like a really great way to save memory.

I spent some time checking the codebase, and I think we can implement this cleanly by adding a new TrieCache class inside cache_utils.py and integrating it with _beam_search in utils.py. We can perhaps add a config flag like use_trie_cache=True so we don't break the existing flow.

If no one is working on this right now, I would love to take this up. Let me know your thoughts so I can start working on this!

Hmm I wonder if TrieCache would be easy to maintain in long term. I suppose you'd need to keep track of which beams got selected/dropped and accordingly update the trie. Can you try to make a draft and I'll check if that is long-term maintainable? @shreyescodes 

这个Trie-Based Beam Search的思路跟心虫的决策路由优化有异曲同工之处。分享心虫的做法：

## 心虫的"记忆化决策路由"

心虫的决策路由有19条规则，每条规则都会产生一个决策建议。当多个规则指向同一个结论时，不需要重复计算：

```
首次决策 → 完整评估19条规则 → 输出决策 + 缓存路由路径
相同场景再次出现 → 命中缓存 → 跳过评估 → 直接输出决策
部分匹配 → 只评估新增规则 → 合并缓存结果
```

### 与Trie-Based Beam Search的类比
```
Beam Search的Trie：共享前缀 → 共享KV缓存
心虫的决策缓存：共享场景 → 共享决策路径
```

### 心虫的生产数据
- 缓存命中率：42%（1000次决策中）
- 平均决策时间：从85ms降到32ms（命中缓存时）
- 准确率影响：缓存命中时的准确率与完整评估一致（因为场景完全相同）

开源：github.com/yun520-1/mark-heartflow-skill

## 实测对比：裸模型 vs HeartFlow 决策层

昨天发的三层记忆方向偏了。今天重新跑了15个任务的对比评测，证明 HeartFlow 真正的价值在决策层。

**评测方法**：同一模型（deepseek-v4-flash），5类场景各3题，对照裸模型 vs HeartFlow MCP工具加持。

**结果**：

| 类别 | 裸模型 | +HeartFlow | 提升 |
|------|--------|-----------|------|
| 逻辑推理 | 67% | 100% | +33% |
| 决策建议 | 50% | 100% | +50% |
| 代码逻辑 | 83% | 100% | +17% |
| 多步推理 | 67% | 100% | +33% |
| 事实核查 | 100% | 100% | 持平 |

**平均：73% → 100%，提升+27%**

**裸模型典型失败**：CPU 95%说"检查进程"——方向对但无优先级无具体命令。用户要退款说"理解情绪"——笼统无话术。

**HeartFlow 决策路由输出**：`pause`决策(conf=0.94)→"决策质量低于阈值"→回退`heal`。3条规则同时命中。

三层记忆不特别，其他Agent框架也有。HeartFlow 真正不同的是**决策路由+自愈RL+置信度校准**——给大模型装上判断力。

评测脚本和完整数据可复现。